Rework UniFi websocket (#100614)

* Rework websocket management

* remove unnecessary fixture

* Remove controller from mock_unifi_websocket

* Mock api.login in reconnect method

* Remove unnecessary edits

* Minor clean up

* Bump aiounifi to v63

* Wait on task cancellation

homeassistant/components/unifi/__init__.py

@@ -52,7 +52,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     if len(hass.data[UNIFI_DOMAIN]) == 1:
         async_setup_services(hass)
 
-    api.start_websocket()
+    controller.start_websocket()
 
     config_entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)

homeassistant/components/unifi/controller.py

@@ -12,7 +12,6 @@ import aiounifi
 from aiounifi.interfaces.api_handlers import ItemEvent
 from aiounifi.models.configuration import Configuration
 from aiounifi.models.device import DeviceSetPoePortModeRequest
-from aiounifi.websocket import WebsocketState
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
@@ -81,7 +80,7 @@ class UniFiController:
         self.config_entry = config_entry
         self.api = api
 
-        api.ws_state_callback = self.async_unifi_ws_state_callback
+        self.ws_task: asyncio.Task | None = None
 
         self.available = True
         self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
@@ -223,23 +222,6 @@ class UniFiController:
         for description in descriptions:
             async_load_entities(description)
 
-    @callback
-    def async_unifi_ws_state_callback(self, state: WebsocketState) -> None:
-        """Handle messages back from UniFi library."""
-        if state == WebsocketState.DISCONNECTED and self.available:
-            LOGGER.warning("Lost connection to UniFi Network")
-
-        if (state == WebsocketState.RUNNING and not self.available) or (
-            state == WebsocketState.DISCONNECTED and self.available
-        ):
-            self.available = state == WebsocketState.RUNNING
-            async_dispatcher_send(self.hass, self.signal_reachable)
-
-            if not self.available:
-                self.hass.loop.call_later(RETRY_TIMER, self.reconnect, True)
-            else:
-                LOGGER.info("Connected to UniFi Network")
-
     @property
     def signal_reachable(self) -> str:
         """Integration specific event to signal a change in connection status."""
@@ -367,6 +349,19 @@ class UniFiController:
         controller.load_config_entry_options()
         async_dispatcher_send(hass, controller.signal_options_update)
 
+    @callback
+    def start_websocket(self) -> None:
+        """Start up connection to websocket."""
+
+        async def _websocket_runner() -> None:
+            """Start websocket."""
+            await self.api.start_websocket()
+            self.available = False
+            async_dispatcher_send(self.hass, self.signal_reachable)
+            self.hass.loop.call_later(RETRY_TIMER, self.reconnect, True)
+
+        self.ws_task = self.hass.loop.create_task(_websocket_runner())
+
     @callback
     def reconnect(self, log: bool = False) -> None:
         """Prepare to reconnect UniFi session."""
@@ -379,7 +374,11 @@ class UniFiController:
         try:
             async with asyncio.timeout(5):
                 await self.api.login()
-                self.api.start_websocket()
+                self.start_websocket()
+
+            if not self.available:
+                self.available = True
+                async_dispatcher_send(self.hass, self.signal_reachable)
 
         except (
             asyncio.TimeoutError,
@@ -395,7 +394,8 @@ class UniFiController:
 
         Used as an argument to EventBus.async_listen_once.
         """
-        self.api.stop_websocket()
+        if self.ws_task is not None:
+            self.ws_task.cancel()
 
     async def async_reset(self) -> bool:
         """Reset this controller to default state.
@@ -403,7 +403,18 @@ class UniFiController:
         Will cancel any scheduled setup retry and will unload
         the config entry.
         """
-        self.api.stop_websocket()
+        if self.ws_task is not None:
+            self.ws_task.cancel()
+
+            _, pending = await asyncio.wait([self.ws_task], timeout=10)
+
+            if pending:
+                LOGGER.warning(
+                    "Unloading %s (%s) config entry. Task %s did not complete in time",
+                    self.config_entry.title,
+                    self.config_entry.domain,
+                    self.ws_task,
+                )
 
         unload_ok = await self.hass.config_entries.async_unload_platforms(
             self.config_entry, PLATFORMS

homeassistant/components/unifi/manifest.json

@@ -8,7 +8,7 @@
   "iot_class": "local_push",
   "loggers": ["aiounifi"],
   "quality_scale": "platinum",
-  "requirements": ["aiounifi==62"],
+  "requirements": ["aiounifi==63"],
   "ssdp": [
     {
       "manufacturer": "Ubiquiti Networks",

requirements_all.txt

@@ -363,7 +363,7 @@ aiosyncthing==0.5.1
 aiotractive==0.5.6
 
 # homeassistant.components.unifi
-aiounifi==62
+aiounifi==63
 
 # homeassistant.components.vlc_telnet
 aiovlc==0.1.0

requirements_test_all.txt

@@ -338,7 +338,7 @@ aiosyncthing==0.5.1
 aiotractive==0.5.6
 
 # homeassistant.components.unifi
-aiounifi==62
+aiounifi==63
 
 # homeassistant.components.vlc_telnet
 aiovlc==0.1.0

tests/components/unifi/conftest.py

@@ -1,47 +1,100 @@
 """Fixtures for UniFi Network methods."""
 from __future__ import annotations
 
+import asyncio
+from datetime import timedelta
 from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
-from aiounifi.websocket import WebsocketSignal, WebsocketState
 import pytest
 
+from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
+from homeassistant.components.unifi.controller import RETRY_TIMER
+from homeassistant.const import 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
+from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.components.unifi.test_controller import DEFAULT_CONFIG_ENTRY_ID
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+
+class WebsocketStateManager(asyncio.Event):
+    """Keep an async event that simules websocket context manager.
+
+    Prepares disconnect and reconnect flows.
+    """
+
+    def __init__(self, hass: HomeAssistant, aioclient_mock: AiohttpClientMocker):
+        """Store hass object and initialize asyncio.Event."""
+        self.hass = hass
+        self.aioclient_mock = aioclient_mock
+        super().__init__()
+
+    async def disconnect(self):
+        """Mark future as done to make 'await self.api.start_websocket' return."""
+        self.set()
+        await self.hass.async_block_till_done()
+
+    async def reconnect(self, fail=False):
+        """Set up new future to make 'await self.api.start_websocket' block.
+
+        Mock api calls done by 'await self.api.login'.
+        Fail will make 'await self.api.start_websocket' return immediately.
+        """
+        controller = self.hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
+        self.aioclient_mock.get(
+            f"https://{controller.host}:1234", status=302
+        )  # Check UniFi OS
+        self.aioclient_mock.post(
+            f"https://{controller.host}:1234/api/login",
+            json={"data": "login successful", "meta": {"rc": "ok"}},
+            headers={"content-type": CONTENT_TYPE_JSON},
+        )
+
+        if not fail:
+            self.clear()
+        new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
+        async_fire_time_changed(self.hass, new_time)
+        await self.hass.async_block_till_done()
 
 
 @pytest.fixture(autouse=True)
-def mock_unifi_websocket():
+def websocket_mock(hass: HomeAssistant, aioclient_mock: AiohttpClientMocker):
-    """No real websocket allowed."""
+    """Mock 'await self.api.start_websocket' in 'UniFiController.start_websocket'."""
-    with patch("aiounifi.controller.WSClient") as mock:
+    websocket_state_manager = WebsocketStateManager(hass, aioclient_mock)
+    with patch("aiounifi.Controller.start_websocket") as ws_mock:
+        ws_mock.side_effect = websocket_state_manager.wait
+        yield websocket_state_manager
 
-        def make_websocket_call(
+
-            *,
+@pytest.fixture(autouse=True)
-            message: MessageKey | None = None,
+def mock_unifi_websocket(hass):
-            data: list[dict] | dict | None = None,
+    """No real websocket allowed."""
-            state: WebsocketState | None = None,
+
-        ):
+    def make_websocket_call(
-            """Generate a websocket call."""
+        *,
-            if data and not message:
+        message: MessageKey | None = None,
-                mock.return_value.data = data
+        data: list[dict] | dict | None = None,
-                mock.call_args[1]["callback"](WebsocketSignal.DATA)
+    ):
-            elif data and message:
+        """Generate a websocket call."""
-                if not isinstance(data, list):
+        controller = hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
-                    data = [data]
+        if data and not message:
-                mock.return_value.data = {
+            controller.api.messages.handler(data)
+        elif data and message:
+            if not isinstance(data, list):
+                data = [data]
+            controller.api.messages.handler(
+                {
                     "meta": {"message": message.value},
                     "data": data,
                 }
-                mock.call_args[1]["callback"](WebsocketSignal.DATA)
+            )
-            elif state:
+        else:
-                mock.return_value.state = state
+            raise NotImplementedError
-                mock.call_args[1]["callback"](WebsocketSignal.CONNECTION_STATE)
-            else:
-                raise NotImplementedError
 
-        yield make_websocket_call
+    return make_websocket_call
 
 
 @pytest.fixture(autouse=True)

tests/components/unifi/test_button.py

@@ -1,7 +1,5 @@
 """UniFi Network button platform tests."""
 
-from aiounifi.websocket import WebsocketState
-
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, ButtonDeviceClass
 from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
 from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE, EntityCategory
@@ -14,7 +12,7 @@ from tests.test_util.aiohttp import AiohttpClientMocker
 
 
 async def test_restart_device_button(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
 ) -> None:
     """Test restarting device button."""
     config_entry = await setup_unifi_integration(
@@ -71,11 +69,9 @@ async def test_restart_device_button(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("button.switch_restart").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("button.switch_restart").state != STATE_UNAVAILABLE

tests/components/unifi/test_controller.py

@@ -6,7 +6,6 @@ from http import HTTPStatus
 from unittest.mock import Mock, patch
 
 import aiounifi
-from aiounifi.websocket import WebsocketState
 import pytest
 
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN
@@ -28,7 +27,7 @@ from homeassistant.components.unifi.const import (
     PLATFORMS,
     UNIFI_WIRELESS_CLIENTS,
 )
-from homeassistant.components.unifi.controller import RETRY_TIMER, get_unifi_controller
+from homeassistant.components.unifi.controller import get_unifi_controller
 from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
 from homeassistant.const import (
     CONF_HOST,
@@ -44,7 +43,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.common import MockConfigEntry
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 DEFAULT_CONFIG_ENTRY_ID = "1"
@@ -365,8 +364,8 @@ async def test_reset_fails(
 async def test_connection_state_signalling(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    mock_unifi_websocket,
     mock_device_registry,
+    websocket_mock,
 ) -> None:
     """Verify connection statesignalling and connection state are working."""
     client = {
@@ -381,21 +380,17 @@ async def test_connection_state_signalling(
     # Controller is connected
     assert hass.states.get("device_tracker.client").state == "home"
 
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
-
     # Controller is disconnected
     assert hass.states.get("device_tracker.client").state == "unavailable"
 
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
-
     # Controller is once again connected
     assert hass.states.get("device_tracker.client").state == "home"
 
 
 async def test_reconnect_mechanism(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
 ) -> None:
     """Verify reconnect prints only on first reconnection try."""
     await setup_unifi_integration(hass, aioclient_mock)
@@ -403,21 +398,13 @@ async def test_reconnect_mechanism(
     aioclient_mock.clear_requests()
     aioclient_mock.get(f"https://{DEFAULT_HOST}:1234/", status=HTTPStatus.BAD_GATEWAY)
 
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
-
     assert aioclient_mock.call_count == 0
 
-    new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
+    await websocket_mock.reconnect(fail=True)
-    async_fire_time_changed(hass, new_time)
-    await hass.async_block_till_done()
-
     assert aioclient_mock.call_count == 1
 
-    new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
+    await websocket_mock.reconnect(fail=True)
-    async_fire_time_changed(hass, new_time)
-    await hass.async_block_till_done()
-
     assert aioclient_mock.call_count == 2
 
 
@@ -431,10 +418,7 @@ async def test_reconnect_mechanism(
     ],
 )
 async def test_reconnect_mechanism_exceptions(
-    hass: HomeAssistant,
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock, exception
-    aioclient_mock: AiohttpClientMocker,
-    mock_unifi_websocket,
-    exception,
 ) -> None:
     """Verify async_reconnect calls expected methods."""
     await setup_unifi_integration(hass, aioclient_mock)
@@ -442,11 +426,9 @@ async def test_reconnect_mechanism_exceptions(
     with patch("aiounifi.Controller.login", side_effect=exception), patch(
         "homeassistant.components.unifi.controller.UniFiController.reconnect"
     ) as mock_reconnect:
-        mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+        await websocket_mock.disconnect()
-        await hass.async_block_till_done()
 
-        new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
+        await websocket_mock.reconnect()
-        async_fire_time_changed(hass, new_time)
         mock_reconnect.assert_called_once()
 
 

tests/components/unifi/test_device_tracker.py

@@ -3,7 +3,6 @@ from datetime import timedelta
 from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
-from aiounifi.websocket import WebsocketState
 from freezegun.api import FrozenDateTimeFactory
 
 from homeassistant import config_entries
@@ -40,8 +39,8 @@ async def test_no_entities(
 async def test_tracked_wireless_clients(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    mock_unifi_websocket,
     mock_device_registry,
+    mock_unifi_websocket,
 ) -> None:
     """Verify tracking of wireless clients."""
     client = {
@@ -402,7 +401,7 @@ async def test_remove_clients(
 async def test_controller_state_change(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    mock_unifi_websocket,
+    websocket_mock,
     mock_device_registry,
 ) -> None:
     """Verify entities state reflect on controller becoming unavailable."""
@@ -443,16 +442,12 @@ async def test_controller_state_change(
     assert hass.states.get("device_tracker.device").state == STATE_HOME
 
     # Controller unavailable
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
-
     assert hass.states.get("device_tracker.client").state == STATE_UNAVAILABLE
     assert hass.states.get("device_tracker.device").state == STATE_UNAVAILABLE
 
     # Controller available
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
-
     assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
     assert hass.states.get("device_tracker.device").state == STATE_HOME
 

tests/components/unifi/test_image.py

@@ -5,7 +5,6 @@ from datetime import timedelta
 from http import HTTPStatus
 
 from aiounifi.models.message import MessageKey
-from aiounifi.websocket import WebsocketState
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.image import DOMAIN as IMAGE_DOMAIN
@@ -65,6 +64,7 @@ async def test_wlan_qr_code(
     hass_client: ClientSessionGenerator,
     snapshot: SnapshotAssertion,
     mock_unifi_websocket,
+    websocket_mock,
 ) -> None:
     """Test the update_clients function when no clients are found."""
     await setup_unifi_integration(hass, aioclient_mock, wlans_response=[WLAN])
@@ -121,13 +121,11 @@ async def test_wlan_qr_code(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("image.ssid_1_qr_code").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE
 
     # WLAN gets disabled

tests/components/unifi/test_sensor.py

@@ -4,7 +4,6 @@ from datetime import datetime, timedelta
 from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
-from aiounifi.websocket import WebsocketState
 import pytest
 
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
@@ -562,7 +561,10 @@ async def test_remove_sensors(
 
 
 async def test_poe_port_switches(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    mock_unifi_websocket,
+    websocket_mock,
 ) -> None:
     """Test the update_items function with some clients."""
     await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1])
@@ -607,16 +609,16 @@ async def test_poe_port_switches(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert (
         hass.states.get("sensor.mock_name_port_1_poe_power").state == STATE_UNAVAILABLE
     )
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
+    assert (
-    assert hass.states.get("sensor.mock_name_port_1_poe_power")
+        hass.states.get("sensor.mock_name_port_1_poe_power").state != STATE_UNAVAILABLE
+    )
 
     # Device gets disabled
     device_1["disabled"] = True
@@ -634,7 +636,10 @@ async def test_poe_port_switches(
 
 
 async def test_wlan_client_sensors(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    mock_unifi_websocket,
+    websocket_mock,
 ) -> None:
     """Verify that WLAN client sensors are working as expected."""
     wireless_client_1 = {
@@ -720,13 +725,11 @@ async def test_wlan_client_sensors(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("sensor.ssid_1").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("sensor.ssid_1").state == "0"
 
     # WLAN gets disabled
@@ -837,7 +840,6 @@ async def test_device_uptime(
     now = datetime(2021, 1, 1, 1, 1, 0, tzinfo=dt_util.UTC)
     with patch("homeassistant.util.dt.now", return_value=now):
         await setup_unifi_integration(hass, aioclient_mock, devices_response=[device])
-
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
     assert hass.states.get("sensor.device_uptime").state == "2021-01-01T01:00:00+00:00"
 
@@ -854,7 +856,6 @@ async def test_device_uptime(
     now = datetime(2021, 1, 1, 1, 1, 4, tzinfo=dt_util.UTC)
     with patch("homeassistant.util.dt.now", return_value=now):
         mock_unifi_websocket(message=MessageKey.DEVICE, data=device)
-        await hass.async_block_till_done()
 
     assert hass.states.get("sensor.device_uptime").state == "2021-01-01T01:00:00+00:00"
 
@@ -865,7 +866,6 @@ async def test_device_uptime(
     now = datetime(2021, 2, 1, 1, 1, 0, tzinfo=dt_util.UTC)
     with patch("homeassistant.util.dt.now", return_value=now):
         mock_unifi_websocket(message=MessageKey.DEVICE, data=device)
-        await hass.async_block_till_done()
 
     assert hass.states.get("sensor.device_uptime").state == "2021-02-01T01:00:00+00:00"
 
@@ -908,5 +908,4 @@ async def test_device_temperature(
     # Verify new event change temperature
     device["general_temperature"] = 60
     mock_unifi_websocket(message=MessageKey.DEVICE, data=device)
-    await hass.async_block_till_done()
     assert hass.states.get("sensor.device_temperature").state == "60"

tests/components/unifi/test_switch.py

@@ -3,7 +3,6 @@ from copy import deepcopy
 from datetime import timedelta
 
 from aiounifi.models.message import MessageKey
-from aiounifi.websocket import WebsocketState
 import pytest
 
 from homeassistant.components.switch import (
@@ -1001,7 +1000,10 @@ async def test_block_switches(
 
 
 async def test_dpi_switches(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    mock_unifi_websocket,
+    websocket_mock,
 ) -> None:
     """Test the update_items function with some clients."""
     await setup_unifi_integration(
@@ -1026,13 +1028,11 @@ async def test_dpi_switches(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.block_media_streaming").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.block_media_streaming").state == STATE_OFF
 
     # Remove app
@@ -1128,6 +1128,7 @@ async def test_outlet_switches(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
     mock_unifi_websocket,
+    websocket_mock,
     entity_id: str,
     test_data: any,
     outlet_index: int,
@@ -1192,13 +1193,11 @@ async def test_outlet_switches(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get(f"switch.{entity_id}").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get(f"switch.{entity_id}").state == STATE_OFF
 
     # Device gets disabled
@@ -1320,7 +1319,10 @@ async def test_option_remove_switches(
 
 
 async def test_poe_port_switches(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    mock_unifi_websocket,
+    websocket_mock,
 ) -> None:
     """Test the update_items function with some clients."""
     config_entry = await setup_unifi_integration(
@@ -1408,13 +1410,11 @@ async def test_poe_port_switches(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_OFF
 
     # Device gets disabled
@@ -1431,7 +1431,10 @@ async def test_poe_port_switches(
 
 
 async def test_wlan_switches(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    mock_unifi_websocket,
+    websocket_mock,
 ) -> None:
     """Test control of UniFi WLAN availability."""
     config_entry = await setup_unifi_integration(
@@ -1488,18 +1491,19 @@ async def test_wlan_switches(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.ssid_1").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.ssid_1").state == STATE_OFF
 
 
 async def test_port_forwarding_switches(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    mock_unifi_websocket,
+    websocket_mock,
 ) -> None:
     """Test control of UniFi port forwarding."""
     _data = {
@@ -1570,13 +1574,11 @@ async def test_port_forwarding_switches(
     # Availability signalling
 
     # Controller disconnects
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.unifi_network_plex").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
     assert hass.states.get("switch.unifi_network_plex").state == STATE_OFF
 
     # Remove entity on deleted message

tests/components/unifi/test_update.py

@@ -2,7 +2,6 @@
 from copy import deepcopy
 
 from aiounifi.models.message import MessageKey
-from aiounifi.websocket import WebsocketState
 from yarl import URL
 
 from homeassistant.components.unifi.const import CONF_SITE_ID
@@ -185,26 +184,18 @@ async def test_install(
 
 
 async def test_controller_state_change(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
 ) -> None:
     """Verify entities state reflect on controller becoming unavailable."""
-    await setup_unifi_integration(
+    await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1])
-        hass,
-        aioclient_mock,
-        devices_response=[DEVICE_1],
-    )
 
     assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1
     assert hass.states.get("update.device_1").state == STATE_ON
 
     # Controller unavailable
-    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await websocket_mock.disconnect()
-    await hass.async_block_till_done()
-
     assert hass.states.get("update.device_1").state == STATE_UNAVAILABLE
 
     # Controller available
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await websocket_mock.reconnect()
-    await hass.async_block_till_done()
-
     assert hass.states.get("update.device_1").state == STATE_ON