Use fixtures to setup UniFi config entries (#118126)

homeassistant/components/unifi/__init__.py

@@ -39,8 +39,6 @@ async def async_setup_entry(
     hass: HomeAssistant, config_entry: UnifiConfigEntry
 ) -> bool:
     """Set up the UniFi Network integration."""
-    hass.data.setdefault(UNIFI_DOMAIN, {})
-
     try:
         api = await get_unifi_api(hass, config_entry.data)
 
@@ -62,7 +60,6 @@ async def async_setup_entry(
     config_entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, hub.shutdown)
     )
-
     return True
 
 

tests/components/unifi/conftest.py

@@ -3,22 +3,265 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Callable
 from datetime import timedelta
+from types import MappingProxyType
+from typing import Any
 from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
 import pytest
 
+from homeassistant.components.unifi.const import CONF_SITE_ID, DOMAIN as UNIFI_DOMAIN
 from homeassistant.components.unifi.hub.websocket import RETRY_TIMER
-from homeassistant.const import CONTENT_TYPE_JSON
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import (
+    CONF_HOST,
+    CONF_PASSWORD,
+    CONF_PORT,
+    CONF_USERNAME,
+    CONF_VERIFY_SSL,
+    CONTENT_TYPE_JSON,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 import homeassistant.util.dt as dt_util
 
 from tests.common import MockConfigEntry, async_fire_time_changed
-from tests.components.unifi.test_hub import DEFAULT_CONFIG_ENTRY_ID
 from tests.test_util.aiohttp import AiohttpClientMocker
 
+DEFAULT_CONFIG_ENTRY_ID = "1"
+DEFAULT_HOST = "1.2.3.4"
+DEFAULT_PORT = 1234
+DEFAULT_SITE = "site_id"
+
+
+@pytest.fixture(autouse=True)
+def mock_discovery():
+    """No real network traffic allowed."""
+    with patch(
+        "homeassistant.components.unifi.config_flow._async_discover_unifi",
+        return_value=None,
+    ) as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_device_registry(hass, device_registry: dr.DeviceRegistry):
+    """Mock device registry."""
+    config_entry = MockConfigEntry(domain="something_else")
+    config_entry.add_to_hass(hass)
+
+    for idx, device in enumerate(
+        (
+            "00:00:00:00:00:01",
+            "00:00:00:00:00:02",
+            "00:00:00:00:00:03",
+            "00:00:00:00:00:04",
+            "00:00:00:00:00:05",
+            "00:00:00:00:00:06",
+            "00:00:00:00:01:01",
+            "00:00:00:00:02:02",
+        )
+    ):
+        device_registry.async_get_or_create(
+            name=f"Device {idx}",
+            config_entry_id=config_entry.entry_id,
+            connections={(dr.CONNECTION_NETWORK_MAC, device)},
+        )
+
+
+# Config entry fixtures
+
+
+@pytest.fixture(name="config_entry")
+def config_entry_fixture(
+    hass: HomeAssistant,
+    config_entry_data: MappingProxyType[str, Any],
+    config_entry_options: MappingProxyType[str, Any],
+) -> ConfigEntry:
+    """Define a config entry fixture."""
+    config_entry = MockConfigEntry(
+        domain=UNIFI_DOMAIN,
+        entry_id="1",
+        unique_id="1",
+        data=config_entry_data,
+        options=config_entry_options,
+        version=1,
+    )
+    config_entry.add_to_hass(hass)
+    return config_entry
+
+
+@pytest.fixture(name="config_entry_data")
+def config_entry_data_fixture() -> MappingProxyType[str, Any]:
+    """Define a config entry data fixture."""
+    return {
+        CONF_HOST: DEFAULT_HOST,
+        CONF_USERNAME: "username",
+        CONF_PASSWORD: "password",
+        CONF_PORT: DEFAULT_PORT,
+        CONF_SITE_ID: DEFAULT_SITE,
+        CONF_VERIFY_SSL: False,
+    }
+
+
+@pytest.fixture(name="config_entry_options")
+def config_entry_options_fixture() -> MappingProxyType[str, Any]:
+    """Define a config entry options fixture."""
+    return {}
+
+
+@pytest.fixture(name="mock_unifi_requests")
+def default_request_fixture(
+    aioclient_mock: AiohttpClientMocker,
+    client_payload: list[dict[str, Any]],
+    clients_all_payload: list[dict[str, Any]],
+    device_payload: list[dict[str, Any]],
+    dpi_app_payload: list[dict[str, Any]],
+    dpi_group_payload: list[dict[str, Any]],
+    port_forward_payload: list[dict[str, Any]],
+    site_payload: list[dict[str, Any]],
+    system_information_payload: list[dict[str, Any]],
+    wlan_payload: list[dict[str, Any]],
+) -> Callable[[str], None]:
+    """Mock default UniFi requests responses."""
+
+    def __mock_default_requests(host: str, site_id: str) -> None:
+        url = f"https://{host}:{DEFAULT_PORT}"
+
+        def mock_get_request(path: str, payload: list[dict[str, Any]]) -> None:
+            aioclient_mock.get(
+                f"{url}{path}",
+                json={"meta": {"rc": "OK"}, "data": payload},
+                headers={"content-type": CONTENT_TYPE_JSON},
+            )
+
+        aioclient_mock.get(url, status=302)  # UniFI OS check
+        aioclient_mock.post(
+            f"{url}/api/login",
+            json={"data": "login successful", "meta": {"rc": "ok"}},
+            headers={"content-type": CONTENT_TYPE_JSON},
+        )
+        mock_get_request("/api/self/sites", site_payload)
+        mock_get_request(f"/api/s/{site_id}/stat/sta", client_payload)
+        mock_get_request(f"/api/s/{site_id}/rest/user", clients_all_payload)
+        mock_get_request(f"/api/s/{site_id}/stat/device", device_payload)
+        mock_get_request(f"/api/s/{site_id}/rest/dpiapp", dpi_app_payload)
+        mock_get_request(f"/api/s/{site_id}/rest/dpigroup", dpi_group_payload)
+        mock_get_request(f"/api/s/{site_id}/rest/portforward", port_forward_payload)
+        mock_get_request(f"/api/s/{site_id}/stat/sysinfo", system_information_payload)
+        mock_get_request(f"/api/s/{site_id}/rest/wlanconf", wlan_payload)
+
+    return __mock_default_requests
+
+
+# Request payload fixtures
+
+
+@pytest.fixture(name="client_payload")
+def client_data_fixture() -> list[dict[str, Any]]:
+    """Client data."""
+    return []
+
+
+@pytest.fixture(name="clients_all_payload")
+def clients_all_data_fixture() -> list[dict[str, Any]]:
+    """Clients all data."""
+    return []
+
+
+@pytest.fixture(name="device_payload")
+def device_data_fixture() -> list[dict[str, Any]]:
+    """Device data."""
+    return []
+
+
+@pytest.fixture(name="dpi_app_payload")
+def dpi_app_data_fixture() -> list[dict[str, Any]]:
+    """DPI app data."""
+    return []
+
+
+@pytest.fixture(name="dpi_group_payload")
+def dpi_group_data_fixture() -> list[dict[str, Any]]:
+    """DPI group data."""
+    return []
+
+
+@pytest.fixture(name="port_forward_payload")
+def port_forward_data_fixture() -> list[dict[str, Any]]:
+    """Port forward data."""
+    return []
+
+
+@pytest.fixture(name="site_payload")
+def site_data_fixture() -> list[dict[str, Any]]:
+    """Site data."""
+    return [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}]
+
+
+@pytest.fixture(name="system_information_payload")
+def system_information_data_fixture() -> list[dict[str, Any]]:
+    """System information data."""
+    return [
+        {
+            "anonymous_controller_id": "24f81231-a456-4c32-abcd-f5612345385f",
+            "build": "atag_7.4.162_21057",
+            "console_display_version": "3.1.15",
+            "hostname": "UDMP",
+            "name": "UDMP",
+            "previous_version": "7.4.156",
+            "timezone": "Europe/Stockholm",
+            "ubnt_device_type": "UDMPRO",
+            "udm_version": "3.0.20.9281",
+            "update_available": False,
+            "update_downloaded": False,
+            "uptime": 1196290,
+            "version": "7.4.162",
+        }
+    ]
+
+
+@pytest.fixture(name="wlan_payload")
+def wlan_data_fixture() -> list[dict[str, Any]]:
+    """WLAN data."""
+    return []
+
+
+@pytest.fixture(name="setup_default_unifi_requests")
+def default_vapix_requests_fixture(
+    config_entry: ConfigEntry,
+    mock_unifi_requests: Callable[[str, str], None],
+) -> None:
+    """Mock default UniFi requests responses."""
+    mock_unifi_requests(config_entry.data[CONF_HOST], config_entry.data[CONF_SITE_ID])
+
+
+@pytest.fixture(name="prepare_config_entry")
+async def prep_config_entry_fixture(
+    hass: HomeAssistant, config_entry: ConfigEntry, setup_default_unifi_requests: None
+) -> Callable[[], ConfigEntry]:
+    """Fixture factory to set up UniFi network integration."""
+
+    async def __mock_setup_config_entry() -> ConfigEntry:
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+        return config_entry
+
+    return __mock_setup_config_entry
+
+
+@pytest.fixture(name="setup_config_entry")
+async def setup_config_entry_fixture(
+    hass: HomeAssistant, prepare_config_entry: Callable[[], ConfigEntry]
+) -> ConfigEntry:
+    """Fixture to set up UniFi network integration."""
+    return await prepare_config_entry()
+
+
+# Websocket fixtures
+
 
 class WebsocketStateManager(asyncio.Event):
     """Keep an async event that simules websocket context manager.
@@ -97,38 +340,3 @@ def mock_unifi_websocket(hass):
             raise NotImplementedError
 
     return make_websocket_call
-
-
-@pytest.fixture(autouse=True)
-def mock_discovery():
-    """No real network traffic allowed."""
-    with patch(
-        "homeassistant.components.unifi.config_flow._async_discover_unifi",
-        return_value=None,
-    ) as mock:
-        yield mock
-
-
-@pytest.fixture
-def mock_device_registry(hass, device_registry: dr.DeviceRegistry):
-    """Mock device registry."""
-    config_entry = MockConfigEntry(domain="something_else")
-    config_entry.add_to_hass(hass)
-
-    for idx, device in enumerate(
-        (
-            "00:00:00:00:00:01",
-            "00:00:00:00:00:02",
-            "00:00:00:00:00:03",
-            "00:00:00:00:00:04",
-            "00:00:00:00:00:05",
-            "00:00:00:00:00:06",
-            "00:00:00:00:01:01",
-            "00:00:00:00:02:02",
-        )
-    ):
-        device_registry.async_get_or_create(
-            name=f"Device {idx}",
-            config_entry_id=config_entry.entry_id,
-            connections={(dr.CONNECTION_NETWORK_MAC, device)},
-        )

tests/components/unifi/test_button.py

@@ -2,6 +2,8 @@
 
 from datetime import timedelta
 
+import pytest
+
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, ButtonDeviceClass
 from homeassistant.components.unifi.const import CONF_SITE_ID
 from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
@@ -17,8 +19,6 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 import homeassistant.util.dt as dt_util
 
-from .test_hub import setup_unifi_integration
-
 from tests.common import async_fire_time_changed
 from tests.test_util.aiohttp import AiohttpClientMocker
 
@@ -60,17 +60,10 @@ WLAN = {
 }
 
 
-async def test_restart_device_button(
+@pytest.mark.parametrize(
-    hass: HomeAssistant,
+    "device_payload",
-    entity_registry: er.EntityRegistry,
+    [
-    aioclient_mock: AiohttpClientMocker,
+        [
-    websocket_mock,
-) -> None:
-    """Test restarting device button."""
-    config_entry = await setup_unifi_integration(
-        hass,
-        aioclient_mock,
-        devices_response=[
             {
                 "board_rev": 3,
                 "device_id": "mock-id",
@@ -83,8 +76,18 @@ async def test_restart_device_button(
                 "type": "usw",
                 "version": "4.0.42.10433",
             }
-        ],
+        ]
-    )
+    ],
+)
+async def test_restart_device_button(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    aioclient_mock: AiohttpClientMocker,
+    setup_config_entry,
+    websocket_mock,
+) -> None:
+    """Test restarting device button."""
+    config_entry = setup_config_entry
     assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 1
 
     ent_reg_entry = entity_registry.async_get("button.switch_restart")
@@ -127,17 +130,10 @@ async def test_restart_device_button(
     assert hass.states.get("button.switch_restart").state != STATE_UNAVAILABLE
 
 
-async def test_power_cycle_poe(
+@pytest.mark.parametrize(
-    hass: HomeAssistant,
+    "device_payload",
-    entity_registry: er.EntityRegistry,
+    [
-    aioclient_mock: AiohttpClientMocker,
+        [
-    websocket_mock,
-) -> None:
-    """Test restarting device button."""
-    config_entry = await setup_unifi_integration(
-        hass,
-        aioclient_mock,
-        devices_response=[
             {
                 "board_rev": 3,
                 "device_id": "mock-id",
@@ -166,8 +162,18 @@ async def test_power_cycle_poe(
                     },
                 ],
             }
-        ],
+        ]
-    )
+    ],
+)
+async def test_power_cycle_poe(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    aioclient_mock: AiohttpClientMocker,
+    setup_config_entry,
+    websocket_mock,
+) -> None:
+    """Test restarting device button."""
+    config_entry = setup_config_entry
     assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 2
 
     ent_reg_entry = entity_registry.async_get("button.switch_port_1_power_cycle")
@@ -214,17 +220,16 @@ async def test_power_cycle_poe(
     )
 
 
+@pytest.mark.parametrize("wlan_payload", [[WLAN]])
 async def test_wlan_regenerate_password(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
+    setup_config_entry,
     websocket_mock,
 ) -> None:
     """Test WLAN regenerate password button."""
-
+    config_entry = setup_config_entry
-    config_entry = await setup_unifi_integration(
-        hass, aioclient_mock, wlans_response=[WLAN]
-    )
     assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 0
 
     button_regenerate_password = "button.ssid_1_regenerate_password"

tests/components/unifi/test_hub.py

@@ -1,9 +1,10 @@
 """Test UniFi Network."""
 
+from collections.abc import Callable
 from copy import deepcopy
 from datetime import timedelta
 from http import HTTPStatus
-from unittest.mock import Mock, patch
+from unittest.mock import patch
 
 import aiounifi
 import pytest
@@ -15,8 +16,6 @@ from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.components.unifi.const import (
     CONF_SITE_ID,
-    CONF_TRACK_CLIENTS,
-    CONF_TRACK_DEVICES,
     DEFAULT_ALLOW_BANDWIDTH_SENSORS,
     DEFAULT_ALLOW_UPTIME_SENSORS,
     DEFAULT_DETECTION_TIME,
@@ -29,6 +28,7 @@ from homeassistant.components.unifi.const import (
 from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
 from homeassistant.components.unifi.hub import get_unifi_api
 from homeassistant.components.update import DOMAIN as UPDATE_DOMAIN
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -39,7 +39,6 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
@@ -239,18 +238,15 @@ async def setup_unifi_integration(
 
 
 async def test_hub_setup(
-    hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
-    aioclient_mock: AiohttpClientMocker,
+    prepare_config_entry: Callable[[], ConfigEntry],
 ) -> None:
     """Successful setup."""
     with patch(
         "homeassistant.config_entries.ConfigEntries.async_forward_entry_setups",
         return_value=True,
     ) as forward_entry_setup:
-        config_entry = await setup_unifi_integration(
+        config_entry = await prepare_config_entry()
-            hass, aioclient_mock, system_information_response=SYSTEM_INFORMATION
-        )
         hub = config_entry.runtime_data
 
     entry = hub.config.entry
@@ -291,109 +287,53 @@ async def test_hub_setup(
     assert device_entry.sw_version == "7.4.162"
 
 
-async def test_hub_not_accessible(hass: HomeAssistant) -> None:
-    """Retry to login gets scheduled when connection fails."""
-    with patch(
-        "homeassistant.components.unifi.hub.get_unifi_api",
-        side_effect=CannotConnect,
-    ):
-        await setup_unifi_integration(hass)
-    assert hass.data[UNIFI_DOMAIN] == {}
-
-
-async def test_hub_trigger_reauth_flow(hass: HomeAssistant) -> None:
-    """Failed authentication trigger a reauthentication flow."""
-    with (
-        patch(
-            "homeassistant.components.unifi.get_unifi_api",
-            side_effect=AuthenticationRequired,
-        ),
-        patch.object(hass.config_entries.flow, "async_init") as mock_flow_init,
-    ):
-        await setup_unifi_integration(hass)
-        mock_flow_init.assert_called_once()
-    assert hass.data[UNIFI_DOMAIN] == {}
-
-
-async def test_hub_unknown_error(hass: HomeAssistant) -> None:
-    """Unknown errors are handled."""
-    with patch(
-        "homeassistant.components.unifi.hub.get_unifi_api",
-        side_effect=Exception,
-    ):
-        await setup_unifi_integration(hass)
-    assert hass.data[UNIFI_DOMAIN] == {}
-
-
-async def test_config_entry_updated(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
-) -> None:
-    """Calling reset when the entry has been setup."""
-    config_entry = await setup_unifi_integration(hass, aioclient_mock)
-    hub = config_entry.runtime_data
-
-    event_call = Mock()
-    unsub = async_dispatcher_connect(hass, hub.signal_options_update, event_call)
-
-    hass.config_entries.async_update_entry(
-        config_entry, options={CONF_TRACK_CLIENTS: False, CONF_TRACK_DEVICES: False}
-    )
-    await hass.async_block_till_done()
-
-    assert config_entry.options[CONF_TRACK_CLIENTS] is False
-    assert config_entry.options[CONF_TRACK_DEVICES] is False
-
-    event_call.assert_called_once()
-
-    unsub()
-
-
 async def test_reset_after_successful_setup(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+    hass: HomeAssistant, setup_config_entry: ConfigEntry
 ) -> None:
     """Calling reset when the entry has been setup."""
-    config_entry = await setup_unifi_integration(hass, aioclient_mock)
+    config_entry = setup_config_entry
-    hub = config_entry.runtime_data
+    assert config_entry.state is ConfigEntryState.LOADED
 
-    result = await hub.async_reset()
+    assert await hass.config_entries.async_unload(config_entry.entry_id)
-    await hass.async_block_till_done()
+    assert config_entry.state is ConfigEntryState.NOT_LOADED
-
-    assert result is True
 
 
 async def test_reset_fails(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+    hass: HomeAssistant, setup_config_entry: ConfigEntry
 ) -> None:
     """Calling reset when the entry has been setup can return false."""
-    config_entry = await setup_unifi_integration(hass, aioclient_mock)
+    config_entry = setup_config_entry
-    hub = config_entry.runtime_data
+    assert config_entry.state is ConfigEntryState.LOADED
 
     with patch(
         "homeassistant.config_entries.ConfigEntries.async_forward_entry_unload",
         return_value=False,
     ):
-        result = await hub.async_reset()
+        assert not await hass.config_entries.async_unload(config_entry.entry_id)
-        await hass.async_block_till_done()
+        assert config_entry.state is ConfigEntryState.LOADED
-
-    assert result is False
 
 
+@pytest.mark.parametrize(
+    "client_payload",
+    [
+        [
+            {
+                "hostname": "client",
+                "ip": "10.0.0.1",
+                "is_wired": True,
+                "last_seen": dt_util.as_timestamp(dt_util.utcnow()),
+                "mac": "00:00:00:00:00:01",
+            },
+        ]
+    ],
+)
 async def test_connection_state_signalling(
     hass: HomeAssistant,
-    aioclient_mock: AiohttpClientMocker,
     mock_device_registry,
+    setup_config_entry: ConfigEntry,
     websocket_mock,
 ) -> None:
     """Verify connection statesignalling and connection state are working."""
-    client = {
-        "hostname": "client",
-        "ip": "10.0.0.1",
-        "is_wired": True,
-        "last_seen": dt_util.as_timestamp(dt_util.utcnow()),
-        "mac": "00:00:00:00:00:01",
-    }
-    await setup_unifi_integration(hass, aioclient_mock, clients_response=[client])
-
     # Controller is connected
     assert hass.states.get("device_tracker.client").state == "home"
 
@@ -407,11 +347,12 @@ async def test_connection_state_signalling(
 
 
 async def test_reconnect_mechanism(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    setup_config_entry: ConfigEntry,
+    websocket_mock,
 ) -> None:
     """Verify reconnect prints only on first reconnection try."""
-    await setup_unifi_integration(hass, aioclient_mock)
-
     aioclient_mock.clear_requests()
     aioclient_mock.get(f"https://{DEFAULT_HOST}:1234/", status=HTTPStatus.BAD_GATEWAY)
 
@@ -435,11 +376,13 @@ async def test_reconnect_mechanism(
     ],
 )
 async def test_reconnect_mechanism_exceptions(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock, exception
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    setup_config_entry: ConfigEntry,
+    websocket_mock,
+    exception,
 ) -> None:
     """Verify async_reconnect calls expected methods."""
-    await setup_unifi_integration(hass, aioclient_mock)
-
     with (
         patch("aiounifi.Controller.login", side_effect=exception),
         patch(
@@ -452,20 +395,6 @@ async def test_reconnect_mechanism_exceptions(
         mock_reconnect.assert_called_once()
 
 
-async def test_get_unifi_api(hass: HomeAssistant) -> None:
-    """Successful call."""
-    with patch("aiounifi.Controller.login", return_value=True):
-        assert await get_unifi_api(hass, ENTRY_CONFIG)
-
-
-async def test_get_unifi_api_verify_ssl_false(hass: HomeAssistant) -> None:
-    """Successful call with verify ssl set to false."""
-    hub_data = dict(ENTRY_CONFIG)
-    hub_data[CONF_VERIFY_SSL] = False
-    with patch("aiounifi.Controller.login", return_value=True):
-        assert await get_unifi_api(hass, hub_data)
-
-
 @pytest.mark.parametrize(
     ("side_effect", "raised_exception"),
     [

tests/components/unifi/test_init.py

@@ -1,9 +1,11 @@
 """Test UniFi Network integration setup process."""
 
+from collections.abc import Callable
 from typing import Any
 from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
+import pytest
 
 from homeassistant.components import unifi
 from homeassistant.components.unifi.const import (
@@ -14,11 +16,12 @@ from homeassistant.components.unifi.const import (
     DOMAIN as UNIFI_DOMAIN,
 )
 from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 from homeassistant.setup import async_setup_component
 
-from .test_hub import DEFAULT_CONFIG_ENTRY_ID, setup_unifi_integration
+from .test_hub import DEFAULT_CONFIG_ENTRY_ID
 
 from tests.common import flush_store
 from tests.test_util.aiohttp import AiohttpClientMocker
@@ -31,18 +34,22 @@ async def test_setup_with_no_config(hass: HomeAssistant) -> None:
     assert UNIFI_DOMAIN not in hass.data
 
 
-async def test_setup_entry_fails_config_entry_not_ready(hass: HomeAssistant) -> None:
+async def test_setup_entry_fails_config_entry_not_ready(
+    hass: HomeAssistant, prepare_config_entry: Callable[[], ConfigEntry]
+) -> None:
     """Failed authentication trigger a reauthentication flow."""
     with patch(
         "homeassistant.components.unifi.get_unifi_api",
         side_effect=CannotConnect,
     ):
-        await setup_unifi_integration(hass)
+        config_entry = await prepare_config_entry()
 
-    assert hass.data[UNIFI_DOMAIN] == {}
+    assert config_entry.state == ConfigEntryState.SETUP_RETRY
 
 
-async def test_setup_entry_fails_trigger_reauth_flow(hass: HomeAssistant) -> None:
+async def test_setup_entry_fails_trigger_reauth_flow(
+    hass: HomeAssistant, prepare_config_entry: Callable[[], ConfigEntry]
+) -> None:
     """Failed authentication trigger a reauthentication flow."""
     with (
         patch(
@@ -51,16 +58,35 @@ async def test_setup_entry_fails_trigger_reauth_flow(hass: HomeAssistant) -> Non
         ),
         patch.object(hass.config_entries.flow, "async_init") as mock_flow_init,
     ):
-        await setup_unifi_integration(hass)
+        config_entry = await prepare_config_entry()
         mock_flow_init.assert_called_once()
 
-    assert hass.data[UNIFI_DOMAIN] == {}
+    assert config_entry.state == ConfigEntryState.SETUP_ERROR
 
 
+@pytest.mark.parametrize(
+    "client_payload",
+    [
+        [
+            {
+                "hostname": "client_1",
+                "ip": "10.0.0.1",
+                "is_wired": False,
+                "mac": "00:00:00:00:00:01",
+            },
+            {
+                "hostname": "client_2",
+                "ip": "10.0.0.2",
+                "is_wired": False,
+                "mac": "00:00:00:00:00:02",
+            },
+        ]
+    ],
+)
 async def test_wireless_clients(
     hass: HomeAssistant,
     hass_storage: dict[str, Any],
-    aioclient_mock: AiohttpClientMocker,
+    prepare_config_entry: Callable[[], ConfigEntry],
 ) -> None:
     """Verify wireless clients class."""
     hass_storage[unifi.STORAGE_KEY] = {
@@ -72,21 +98,7 @@ async def test_wireless_clients(
         },
     }
 
-    client_1 = {
+    await prepare_config_entry()
-        "hostname": "client_1",
-        "ip": "10.0.0.1",
-        "is_wired": False,
-        "mac": "00:00:00:00:00:01",
-    }
-    client_2 = {
-        "hostname": "client_2",
-        "ip": "10.0.0.2",
-        "is_wired": False,
-        "mac": "00:00:00:00:00:02",
-    }
-    await setup_unifi_integration(
-        hass, aioclient_mock, clients_response=[client_1, client_2]
-    )
     await flush_store(hass.data[unifi.UNIFI_WIRELESS_CLIENTS]._store)
 
     assert sorted(hass_storage[unifi.STORAGE_KEY]["data"]["wireless_clients"]) == [
@@ -96,98 +108,113 @@ async def test_wireless_clients(
     ]
 
 
+@pytest.mark.parametrize(
+    "client_payload",
+    [
+        [
+            {
+                "hostname": "Wired client",
+                "is_wired": True,
+                "mac": "00:00:00:00:00:01",
+                "oui": "Producer",
+                "wired-rx_bytes": 1234000000,
+                "wired-tx_bytes": 5678000000,
+                "uptime": 1600094505,
+            },
+            {
+                "is_wired": False,
+                "mac": "00:00:00:00:00:02",
+                "name": "Wireless client",
+                "oui": "Producer",
+                "rx_bytes": 2345000000,
+                "tx_bytes": 6789000000,
+                "uptime": 60,
+            },
+        ]
+    ],
+)
+@pytest.mark.parametrize(
+    "device_payload",
+    [
+        [
+            {
+                "board_rev": 3,
+                "device_id": "mock-id",
+                "has_fan": True,
+                "fan_level": 0,
+                "ip": "10.0.1.1",
+                "last_seen": 1562600145,
+                "mac": "00:00:00:00:01:01",
+                "model": "US16P150",
+                "name": "Device 1",
+                "next_interval": 20,
+                "overheating": True,
+                "state": 1,
+                "type": "usw",
+                "upgradable": True,
+                "version": "4.0.42.10433",
+            }
+        ]
+    ],
+)
+@pytest.mark.parametrize(
+    "config_entry_options",
+    [
+        {
+            CONF_ALLOW_BANDWIDTH_SENSORS: True,
+            CONF_ALLOW_UPTIME_SENSORS: True,
+            CONF_TRACK_CLIENTS: True,
+            CONF_TRACK_DEVICES: True,
+        }
+    ],
+)
 async def test_remove_config_entry_device(
     hass: HomeAssistant,
     hass_storage: dict[str, Any],
     aioclient_mock: AiohttpClientMocker,
     device_registry: dr.DeviceRegistry,
+    prepare_config_entry: Callable[[], ConfigEntry],
+    client_payload: list[dict[str, Any]],
+    device_payload: list[dict[str, Any]],
     mock_unifi_websocket,
     hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Verify removing a device manually."""
-    client_1 = {
+    config_entry = await prepare_config_entry()
-        "hostname": "Wired client",
-        "is_wired": True,
-        "mac": "00:00:00:00:00:01",
-        "oui": "Producer",
-        "wired-rx_bytes": 1234000000,
-        "wired-tx_bytes": 5678000000,
-        "uptime": 1600094505,
-    }
-    client_2 = {
-        "is_wired": False,
-        "mac": "00:00:00:00:00:02",
-        "name": "Wireless client",
-        "oui": "Producer",
-        "rx_bytes": 2345000000,
-        "tx_bytes": 6789000000,
-        "uptime": 60,
-    }
-    device_1 = {
-        "board_rev": 3,
-        "device_id": "mock-id",
-        "has_fan": True,
-        "fan_level": 0,
-        "ip": "10.0.1.1",
-        "last_seen": 1562600145,
-        "mac": "00:00:00:00:01:01",
-        "model": "US16P150",
-        "name": "Device 1",
-        "next_interval": 20,
-        "overheating": True,
-        "state": 1,
-        "type": "usw",
-        "upgradable": True,
-        "version": "4.0.42.10433",
-    }
-    options = {
-        CONF_ALLOW_BANDWIDTH_SENSORS: True,
-        CONF_ALLOW_UPTIME_SENSORS: True,
-        CONF_TRACK_CLIENTS: True,
-        CONF_TRACK_DEVICES: True,
-    }
-
-    config_entry = await setup_unifi_integration(
-        hass,
-        aioclient_mock,
-        options=options,
-        clients_response=[client_1, client_2],
-        devices_response=[device_1],
-    )
 
     assert await async_setup_component(hass, "config", {})
     ws_client = await hass_ws_client(hass)
 
     # Try to remove an active client from UI: not allowed
     device_entry = device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, client_1["mac"])}
+        connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])}
     )
     response = await ws_client.remove_device(device_entry.id, config_entry.entry_id)
     assert not response["success"]
     assert device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, client_1["mac"])}
+        connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])}
     )
 
     # Try to remove an active device from UI: not allowed
     device_entry = device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, device_1["mac"])}
+        connections={(dr.CONNECTION_NETWORK_MAC, device_payload[0]["mac"])}
     )
     response = await ws_client.remove_device(device_entry.id, config_entry.entry_id)
     assert not response["success"]
     assert device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, device_1["mac"])}
+        connections={(dr.CONNECTION_NETWORK_MAC, device_payload[0]["mac"])}
     )
 
     # Remove a client from Unifi API
-    mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=[client_2])
+    mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=[client_payload[1]])
     await hass.async_block_till_done()
 
     # Try to remove an inactive client from UI: allowed
     device_entry = device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, client_2["mac"])}
+        connections={(dr.CONNECTION_NETWORK_MAC, client_payload[1]["mac"])}
     )
     response = await ws_client.remove_device(device_entry.id, config_entry.entry_id)
     assert response["success"]
     assert not device_registry.async_get_device(
-        connections={(dr.CONNECTION_NETWORK_MAC, client_2["mac"])}
+        connections={(dr.CONNECTION_NETWORK_MAC, client_payload[1]["mac"])}
     )