From 642a06b0f087867a42462b592d8c92b347795bed Mon Sep 17 00:00:00 2001
From: Arie Catsman <120491684+catsmanac@users.noreply.github.com>
Date: Mon, 27 Jan 2025 00:38:45 +0100
Subject: [PATCH] Optimize enphase_envoy test integration setup. (#136572)

---
 tests/components/enphase_envoy/__init__.py  | 12 +++++++++---
 tests/components/enphase_envoy/test_init.py | 20 +-------------------
 2 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/tests/components/enphase_envoy/__init__.py b/tests/components/enphase_envoy/__init__.py
index f69ab8e44f2..f5381eda2a7 100644
--- a/tests/components/enphase_envoy/__init__.py
+++ b/tests/components/enphase_envoy/__init__.py
@@ -1,13 +1,19 @@
 """Tests for the Enphase Envoy integration."""
 
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry
 
 
-async def setup_integration(hass: HomeAssistant, config_entry: MockConfigEntry) -> None:
-    """Fixture for setting up the component."""
+async def setup_integration(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    expected_state: ConfigEntryState = ConfigEntryState.LOADED,
+) -> None:
+    """Fixture for setting up the component and testing expected state."""
     config_entry.add_to_hass(hass)
 
     await hass.config_entries.async_setup(config_entry.entry_id)
-    await hass.async_block_till_done()
+    await hass.async_block_till_done(wait_background_tasks=True)
+    assert config_entry.state is expected_state
diff --git a/tests/components/enphase_envoy/test_init.py b/tests/components/enphase_envoy/test_init.py
index 620bd654aca..93a150cfc5c 100644
--- a/tests/components/enphase_envoy/test_init.py
+++ b/tests/components/enphase_envoy/test_init.py
@@ -52,8 +52,6 @@ async def test_with_pre_v7_firmware(
     )
     await setup_integration(hass, config_entry)
 
-    assert config_entry.state is ConfigEntryState.LOADED
-
     assert (entity_state := hass.states.get("sensor.inverter_1"))
     assert entity_state.state == "1"
 
@@ -84,8 +82,6 @@ async def test_token_in_config_file(
     )
     mock_envoy.auth = EnvoyTokenAuth("127.0.0.1", token=token, envoy_serial="1234")
     await setup_integration(hass, entry)
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert entry.state is ConfigEntryState.LOADED
 
     assert (entity_state := hass.states.get("sensor.inverter_1"))
     assert entity_state.state == "1"
@@ -129,8 +125,6 @@ async def test_expired_token_in_config(
         cloud_password="test_password",
     )
     await setup_integration(hass, entry)
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert entry.state is ConfigEntryState.LOADED
 
     assert (entity_state := hass.states.get("sensor.inverter_1"))
     assert entity_state.state == "1"
@@ -230,9 +224,6 @@ async def test_coordinator_token_refresh_error(
     ):
         await setup_integration(hass, entry)
 
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert entry.state is ConfigEntryState.LOADED
-
     assert (entity_state := hass.states.get("sensor.inverter_1"))
     assert entity_state.state == "1"
 
@@ -255,7 +246,6 @@ async def test_config_no_unique_id(
         },
     )
     await setup_integration(hass, entry)
-    assert entry.state is ConfigEntryState.LOADED
     assert entry.unique_id == mock_envoy.serial_number
 
 
@@ -276,8 +266,7 @@ async def test_config_different_unique_id(
             CONF_PASSWORD: "test-password",
         },
     )
-    await setup_integration(hass, entry)
-    assert entry.state is ConfigEntryState.SETUP_RETRY
+    await setup_integration(hass, entry, expected_state=ConfigEntryState.SETUP_RETRY)
 
 
 @pytest.mark.parametrize(
@@ -298,7 +287,6 @@ async def test_remove_config_entry_device(
     """Test removing enphase_envoy config entry device."""
     assert await async_setup_component(hass, "config", {})
     await setup_integration(hass, config_entry)
-    assert config_entry.state is ConfigEntryState.LOADED
 
     # use client to send remove_device command
     hass_client = await hass_ws_client(hass)
@@ -349,8 +337,6 @@ async def test_option_change_reload(
 ) -> None:
     """Test options change will reload entity."""
     await setup_integration(hass, config_entry)
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert config_entry.state is ConfigEntryState.LOADED
     # By default neither option is available
     assert config_entry.options == {}
 
@@ -403,8 +389,6 @@ async def test_coordinator_firmware_refresh(
 ) -> None:
     """Test coordinator scheduled firmware check."""
     await setup_integration(hass, config_entry)
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert config_entry.state is ConfigEntryState.LOADED
 
     # Move time to next firmware check moment
     # SCAN_INTERVAL is patched to 1 day to disable it's firmware detection
@@ -447,8 +431,6 @@ async def test_coordinator_firmware_refresh_with_envoy_error(
 ) -> None:
     """Test coordinator scheduled firmware check."""
     await setup_integration(hass, config_entry)
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert config_entry.state is ConfigEntryState.LOADED
 
     caplog.set_level(logging.DEBUG)
     logging.getLogger("homeassistant.components.enphase_envoy.coordinator").setLevel(