From 01b58549685da5e9f4c25f1d4817b2f9ab6afba4 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Wed, 27 Sep 2023 10:56:24 +0200 Subject: [PATCH] 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 | 2 +- homeassistant/components/unifi/controller.py | 55 +++++---- homeassistant/components/unifi/manifest.json | 2 +- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/unifi/conftest.py | 105 +++++++++++++----- tests/components/unifi/test_button.py | 10 +- tests/components/unifi/test_controller.py | 42 ++----- tests/components/unifi/test_device_tracker.py | 13 +-- tests/components/unifi/test_image.py | 8 +- tests/components/unifi/test_sensor.py | 31 +++--- tests/components/unifi/test_switch.py | 52 ++++----- tests/components/unifi/test_update.py | 17 +-- 13 files changed, 184 insertions(+), 157 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 0bde41ac611..4337899a50f 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/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) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 9f965b424ff..620b928176e 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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 diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index 8734fd7dce5..7673402aaac 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/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", diff --git a/requirements_all.txt b/requirements_all.txt index 45bdb08b7fb..013c02c2d2a 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index e78e1d2658d..478db5bffeb 100644 --- a/requirements_test_all.txt +++ b/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 diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index ca0c855d1ab..d48ff613902 100644 --- a/tests/components/unifi/conftest.py +++ b/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(): - """No real websocket allowed.""" - with patch("aiounifi.controller.WSClient") as mock: +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 - 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 = { + +@pytest.fixture(autouse=True) +def mock_unifi_websocket(hass): + """No real websocket allowed.""" + + 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 + ) + else: + raise NotImplementedError - yield make_websocket_call + return make_websocket_call @pytest.fixture(autouse=True) diff --git a/tests/components/unifi/test_button.py b/tests/components/unifi/test_button.py index 0c6ac38739e..30a1b3e08ff 100644 --- a/tests/components/unifi/test_button.py +++ b/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 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 diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index f4738862aef..93b39d2fdf2 100644 --- a/tests/components/unifi/test_controller.py +++ b/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 hass.async_block_till_done() - + await websocket_mock.disconnect() # Controller is disconnected assert hass.states.get("device_tracker.client").state == "unavailable" - mock_unifi_websocket(state=WebsocketState.RUNNING) - await hass.async_block_till_done() - + await websocket_mock.reconnect() # 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 hass.async_block_till_done() - + await websocket_mock.disconnect() assert aioclient_mock.call_count == 0 - new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER) - async_fire_time_changed(hass, new_time) - await hass.async_block_till_done() - + await websocket_mock.reconnect(fail=True) assert aioclient_mock.call_count == 1 - new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER) - async_fire_time_changed(hass, new_time) - await hass.async_block_till_done() - + await websocket_mock.reconnect(fail=True) assert aioclient_mock.call_count == 2 @@ -431,10 +418,7 @@ async def test_reconnect_mechanism( ], ) async def test_reconnect_mechanism_exceptions( - hass: HomeAssistant, - aioclient_mock: AiohttpClientMocker, - mock_unifi_websocket, - exception, + hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock, 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 hass.async_block_till_done() + await websocket_mock.disconnect() - new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER) - async_fire_time_changed(hass, new_time) + await websocket_mock.reconnect() mock_reconnect.assert_called_once() diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 99874b3a949..2680a357d77 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/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 hass.async_block_till_done() - + await websocket_mock.disconnect() 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 hass.async_block_till_done() - + await websocket_mock.reconnect() assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME assert hass.states.get("device_tracker.device").state == STATE_HOME diff --git a/tests/components/unifi/test_image.py b/tests/components/unifi/test_image.py index 38a8cef43c1..92879f5ad14 100644 --- a/tests/components/unifi/test_image.py +++ b/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 hass.async_block_till_done() + await websocket_mock.disconnect() assert hass.states.get("image.ssid_1_qr_code").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("image.ssid_1_qr_code").state != STATE_UNAVAILABLE # WLAN gets disabled diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 7b6a3bc1edc..b652c38abdb 100644 --- a/tests/components/unifi/test_sensor.py +++ b/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 hass.async_block_till_done() + await websocket_mock.disconnect() assert ( hass.states.get("sensor.mock_name_port_1_poe_power").state == STATE_UNAVAILABLE ) # Controller reconnects - mock_unifi_websocket(state=WebsocketState.RUNNING) - await hass.async_block_till_done() - assert hass.states.get("sensor.mock_name_port_1_poe_power") + await websocket_mock.reconnect() + assert ( + 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 hass.async_block_till_done() + await websocket_mock.disconnect() assert hass.states.get("sensor.ssid_1").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("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" diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 8e536119291..a08cf0be688 100644 --- a/tests/components/unifi/test_switch.py +++ b/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 hass.async_block_till_done() + await websocket_mock.disconnect() assert hass.states.get("switch.block_media_streaming").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("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 hass.async_block_till_done() + await websocket_mock.disconnect() assert hass.states.get(f"switch.{entity_id}").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(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 hass.async_block_till_done() + await websocket_mock.disconnect() assert hass.states.get("switch.mock_name_port_1_poe").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("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 hass.async_block_till_done() + await websocket_mock.disconnect() assert hass.states.get("switch.ssid_1").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("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 hass.async_block_till_done() + await websocket_mock.disconnect() assert hass.states.get("switch.unifi_network_plex").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("switch.unifi_network_plex").state == STATE_OFF # Remove entity on deleted message diff --git a/tests/components/unifi/test_update.py b/tests/components/unifi/test_update.py index e59eca371d6..4f7a3dfe11d 100644 --- a/tests/components/unifi/test_update.py +++ b/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( - hass, - aioclient_mock, - devices_response=[DEVICE_1], - ) + await setup_unifi_integration(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 hass.async_block_till_done() - + await websocket_mock.disconnect() assert hass.states.get("update.device_1").state == STATE_UNAVAILABLE # Controller available - mock_unifi_websocket(state=WebsocketState.RUNNING) - await hass.async_block_till_done() - + await websocket_mock.reconnect() assert hass.states.get("update.device_1").state == STATE_ON