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 @@
 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 @@ def __init__(
         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 @@ def async_options_updated() -> None:
         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 @@ async def async_config_entry_updated(
         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 @@ async def async_reconnect(self) -> None:
         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,15 +394,27 @@ def shutdown(self, event: Event) -> None:
 
         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.
 
         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 websocket_mock(hass: HomeAssistant, aioclient_mock: AiohttpClientMocker):
+    """Mock 'await self.api.start_websocket' in 'UniFiController.start_websocket'."""
+    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
 
 
 @pytest.fixture(autouse=True)
-def mock_unifi_websocket():
+def mock_unifi_websocket(hass):
     """No real websocket allowed."""
-    with patch("aiounifi.controller.WSClient") as mock:
-
-        def make_websocket_call(
-            *,
-            message: MessageKey | None = None,
-            data: list[dict] | dict | None = None,
-            state: WebsocketState | None = None,
-        ):
-            """Generate a websocket call."""
-            if data and not message:
-                mock.return_value.data = data
-                mock.call_args[1]["callback"](WebsocketSignal.DATA)
-            elif data and message:
-                if not isinstance(data, list):
-                    data = [data]
-                mock.return_value.data = {
+
+    def make_websocket_call(
+        *,
+        message: MessageKey | None = None,
+        data: list[dict] | dict | None = None,
+    ):
+        """Generate a websocket call."""
+        controller = hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
+        if data and not message:
+            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:
-                mock.return_value.state = state
-                mock.call_args[1]["callback"](WebsocketSignal.CONNECTION_STATE)
-            else:
-                raise NotImplementedError
-
-        yield make_websocket_call
+            )
+        else:
+            raise NotImplementedError
+
+    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 @@
 
 
 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 hass.async_block_till_done()
+    await websocket_mock.disconnect()
     assert hass.states.get("button.switch_restart").state == STATE_UNAVAILABLE
 
     # Controller reconnects
-    mock_unifi_websocket(state=WebsocketState.RUNNING)
-    await hass.async_block_till_done()
+    await websocket_mock.reconnect()
     assert hass.states.get("button.switch_restart").state != STATE_UNAVAILABLE