From 4feb3740dc9e29b2c4a6a55c2c58f55c17c18e54 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Mon, 14 Sep 2026 17:29:26 +0200 Subject: [PATCH] Unifi common data coordinator (#181910) --- homeassistant/components/unifi/coordinator.py | 38 ++- homeassistant/components/unifi/entity.py | 62 +++- .../components/unifi/hub/entity_loader.py | 53 +-- homeassistant/components/unifi/light.py | 2 + homeassistant/components/unifi/switch.py | 12 +- tests/components/unifi/test_hub.py | 316 +++++++++++++++++- tests/components/unifi/test_switch.py | 28 +- 7 files changed, 451 insertions(+), 60 deletions(-) diff --git a/homeassistant/components/unifi/coordinator.py b/homeassistant/components/unifi/coordinator.py index d260b2398754..cca4293d7566 100644 --- a/homeassistant/components/unifi/coordinator.py +++ b/homeassistant/components/unifi/coordinator.py @@ -3,8 +3,9 @@ from datetime import timedelta from typing import TYPE_CHECKING, override -from aiounifi.interfaces.api_handlers import APIHandler +from aiounifi.interfaces.api_handlers import APIHandler, ItemEvent +from homeassistant.core import callback from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from .const import LOGGER @@ -15,8 +16,10 @@ if TYPE_CHECKING: POLL_INTERVAL = timedelta(seconds=10) -class UnifiDataUpdateCoordinator[HandlerT: APIHandler](DataUpdateCoordinator[None]): - """Coordinator managing polling for a single UniFi API data source.""" +class UnifiDataUpdateCoordinator[HandlerT: APIHandler]( + DataUpdateCoordinator[tuple[ItemEvent, str] | None] +): + """Coordinator managing websocket or polling updates for a UniFi API handler.""" def __init__( self, @@ -24,15 +27,18 @@ class UnifiDataUpdateCoordinator[HandlerT: APIHandler](DataUpdateCoordinator[Non handler: HandlerT, ) -> None: """Initialize coordinator.""" + supports_websocket = bool(handler.process_messages or handler.remove_messages) super().__init__( hub.hass, LOGGER, name=f"UniFi {type(handler).__name__}", config_entry=hub.config.entry, - update_interval=POLL_INTERVAL, + update_interval=None if supports_websocket else POLL_INTERVAL, ) self._handler = handler + hub.config.entry.async_on_unload(handler.subscribe(self._async_handle_update)) + @property def handler(self) -> HandlerT: """Return the aiounifi handler managed by this coordinator.""" @@ -42,3 +48,27 @@ class UnifiDataUpdateCoordinator[HandlerT: APIHandler](DataUpdateCoordinator[Non async def _async_update_data(self) -> None: """Update data from the API handler.""" await self._handler.update() + + @callback + def _async_handle_update(self, event: ItemEvent, obj_id: str) -> None: + """Notify listeners which object changed on a websocket update.""" + self.async_set_updated_data((event, obj_id)) + + @callback + @override + def async_update_listeners(self) -> None: + """Notify listeners for the changed object or a polling refresh.""" + data = self.data + changed_obj_id = data[1] if data is not None else None + for update_callback, context in list(self._listeners.values()): + if changed_obj_id is not None and isinstance(context, tuple): + if changed_obj_id not in context: + continue + try: + update_callback() + except Exception: + self.logger.exception( + "Unexpected error updating listener %s for %s", + id(update_callback), + self.name, + ) diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index 114468047319..bacf4748f67b 100644 --- a/homeassistant/components/unifi/entity.py +++ b/homeassistant/components/unifi/entity.py @@ -26,6 +26,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import Entity, EntityDescription from .const import ATTR_MANUFACTURER, DOMAIN +from .coordinator import UnifiDataUpdateCoordinator if TYPE_CHECKING: from .hub import UnifiHub @@ -136,6 +137,7 @@ class UnifiEntity[HandlerT: APIHandler, ItemT: ApiItem](Entity): """Representation of a UniFi entity.""" entity_description: UnifiEntityDescription[HandlerT, ItemT] + coordinator: UnifiDataUpdateCoordinator[HandlerT] _attr_unique_id: str def __init__( @@ -149,6 +151,9 @@ class UnifiEntity[HandlerT: APIHandler, ItemT: ApiItem](Entity): self.hub = hub self.api = hub.api self.entity_description = description + self.coordinator = hub.entity_loader.get_data_update_coordinator( + description.api_handler_fn(self.api) + ) hub.entity_loader.known_objects.add((description.key, obj_id)) @@ -172,7 +177,6 @@ class UnifiEntity[HandlerT: APIHandler, ItemT: ApiItem](Entity): async def async_added_to_hass(self) -> None: """Register callbacks.""" description = self.entity_description - handler = description.api_handler_fn(self.api) @callback def unregister_object() -> None: @@ -183,11 +187,11 @@ class UnifiEntity[HandlerT: APIHandler, ItemT: ApiItem](Entity): self.async_on_remove(unregister_object) - # New data from handler + # New data from coordinator self.async_on_remove( - handler.subscribe( - self.async_signalling_callback, - id_filter=self._obj_id, + self.coordinator.async_add_listener( + self._async_coordinator_updated, + context=(self._obj_id, self._obj_id.partition("_")[0]), ) ) @@ -219,10 +223,29 @@ class UnifiEntity[HandlerT: APIHandler, ItemT: ApiItem](Entity): ) @callback - def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: - """Update the entity state.""" - if event is ItemEvent.DELETED and obj_id == self._obj_id: - self.hass.async_create_task(self.remove_item({obj_id})) + def _async_coordinator_updated(self) -> None: + """Skip coordinator updates that changed a different object.""" + coordinator_data = self.coordinator.data + if coordinator_data is None: + event = ItemEvent.CHANGED + changed_obj_id = None + else: + event, changed_obj_id = coordinator_data + + own_obj_id = self._obj_id.partition("_")[0] + if changed_obj_id is not None and changed_obj_id not in ( + self._obj_id, + own_obj_id, + ): + return + self._async_process_update(event) + + @callback + def _async_process_update(self, event: ItemEvent = ItemEvent.CHANGED) -> None: + """Update the entity state from the handler.""" + handler = self.entity_description.api_handler_fn(self.api) + if self._obj_id not in handler: + self.hass.async_create_task(self.remove_item({self._obj_id})) return description = self.entity_description @@ -230,10 +253,22 @@ class UnifiEntity[HandlerT: APIHandler, ItemT: ApiItem](Entity): self.hass.async_create_task(self.remove_item({self._obj_id})) return - self._attr_available = description.available_fn(self.hub, self._obj_id) - self.async_update_state(event, obj_id) + self._attr_available = ( + description.available_fn(self.hub, self._obj_id) + and self.coordinator.last_update_success + ) + self.async_update_state(event, self._obj_id) self.async_write_ha_state() + @callback + def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: + """Update the entity state from a handler event.""" + if event is ItemEvent.DELETED and obj_id == self._obj_id: + self.hass.async_create_task(self.remove_item({obj_id})) + return + + self._async_process_update(event) + @callback def async_signal_reachable_callback(self) -> None: """Call when hub connection state change.""" @@ -258,6 +293,11 @@ class UnifiEntity[HandlerT: APIHandler, ItemT: ApiItem](Entity): """Update state if polling is configured.""" self.async_update_state(ItemEvent.CHANGED, self._obj_id) + async def async_refresh_after_control(self) -> None: + """Refresh handler data after a control call when polling.""" + if self.coordinator.update_interval is not None: + await self.coordinator.async_refresh() + @callback def async_initiate_state(self) -> None: """Initiate entity state. diff --git a/homeassistant/components/unifi/hub/entity_loader.py b/homeassistant/components/unifi/hub/entity_loader.py index 9e70d51cc70b..e38db332d659 100644 --- a/homeassistant/components/unifi/hub/entity_loader.py +++ b/homeassistant/components/unifi/hub/entity_loader.py @@ -11,6 +11,7 @@ from functools import partial from typing import TYPE_CHECKING, Any from aiounifi.interfaces.api_handlers import APIHandler, ItemEvent +from aiounifi.models.api import ApiItem from aiounifi.models.client import Client from homeassistant.const import Platform @@ -36,33 +37,43 @@ class UnifiEntityLoader: def __init__(self, hub: UnifiHub) -> None: """Initialize the UniFi entity loader.""" self.hub = hub - self.api_updaters = ( - hub.api.clients.update, + self._startup_only_api_updaters = ( hub.api.clients_all.update, - hub.api.devices.update, - hub.api.dpi_apps.update, - hub.api.dpi_groups.update, - hub.api.port_forwarding.update, hub.api.sites.update, hub.api.system_information.update, - hub.api.firewall_policies.update, - hub.api.wlans.update, ) self.wireless_clients = hub.hass.data[UNIFI_WIRELESS_CLIENTS] - self._polling_coordinators: dict[int, UnifiDataUpdateCoordinator] = { + self._data_coordinators: dict[int, UnifiDataUpdateCoordinator[Any]] = { + id(hub.api.clients): UnifiDataUpdateCoordinator(hub, hub.api.clients), + id(hub.api.devices): UnifiDataUpdateCoordinator(hub, hub.api.devices), + id(hub.api.dpi_apps): UnifiDataUpdateCoordinator(hub, hub.api.dpi_apps), + id(hub.api.dpi_groups): UnifiDataUpdateCoordinator(hub, hub.api.dpi_groups), + id(hub.api.firewall_policies): UnifiDataUpdateCoordinator( + hub, hub.api.firewall_policies + ), id(hub.api.object_oriented_network_configs): UnifiDataUpdateCoordinator( hub, hub.api.object_oriented_network_configs ), + id(hub.api.port_forwarding): UnifiDataUpdateCoordinator( + hub, hub.api.port_forwarding + ), id(hub.api.traffic_rules): UnifiDataUpdateCoordinator( hub, hub.api.traffic_rules ), id(hub.api.traffic_routes): UnifiDataUpdateCoordinator( hub, hub.api.traffic_routes ), + id(hub.api.wlans): UnifiDataUpdateCoordinator(hub, hub.api.wlans), } - for coordinator in self._polling_coordinators.values(): - coordinator.async_add_listener(lambda: None) + self._data_coordinator_aliases: dict[int, int] = { + id(hub.api.outlets): id(hub.api.devices), + id(hub.api.ports): id(hub.api.devices), + } + for coordinator in self._data_coordinators.values(): + self.hub.config.entry.async_on_unload( + coordinator.async_add_listener(lambda: None) + ) self.platforms: list[ tuple[ @@ -79,11 +90,11 @@ class UnifiEntityLoader: async def initialize(self) -> None: """Initialize API data and extra client support.""" await asyncio.gather( - self._refresh_api_data(), + self._refresh_data(self._startup_only_api_updaters), self._refresh_data( [ coordinator.async_refresh - for coordinator in self._polling_coordinators.values() + for coordinator in self._data_coordinators.values() ] ), ) @@ -101,10 +112,6 @@ class UnifiEntityLoader: if result is not None: LOGGER.warning("Exception on update %s", result) - async def _refresh_api_data(self) -> None: - """Refresh API data from network application.""" - await self._refresh_data(self.api_updaters) - @callback def _restore_inactive_clients(self) -> None: """Restore recently seen inactive clients and prune stale ones. @@ -214,11 +221,13 @@ class UnifiEntityLoader: ) @callback - def get_data_update_coordinator( - self, handler: APIHandler - ) -> UnifiDataUpdateCoordinator | None: - """Return the polling coordinator for a handler, if available.""" - return self._polling_coordinators.get(id(handler)) + def get_data_update_coordinator[HandlerT: APIHandler[ApiItem]]( + self, handler: HandlerT + ) -> UnifiDataUpdateCoordinator[HandlerT]: + """Return the data coordinator for a handler.""" + handler_id = id(handler) + resolved_handler_id = self._data_coordinator_aliases.get(handler_id, handler_id) + return self._data_coordinators[resolved_handler_id] @callback def _load_entities( diff --git a/homeassistant/components/unifi/light.py b/homeassistant/components/unifi/light.py index 4262e5bb854c..f279edb2e561 100644 --- a/homeassistant/components/unifi/light.py +++ b/homeassistant/components/unifi/light.py @@ -184,6 +184,7 @@ class UnifiLightEntity[HandlerT: APIHandler, ApiItemT: ApiItem]( translation_domain=DOMAIN, translation_key="action_request_failed", ) from err + await self.async_refresh_after_control() @override async def async_turn_off(self, **kwargs: Any) -> None: @@ -197,6 +198,7 @@ class UnifiLightEntity[HandlerT: APIHandler, ApiItemT: ApiItem]( translation_domain=DOMAIN, translation_key="action_request_failed", ) from err + await self.async_refresh_after_control() @callback @override diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 4694b757f26b..f0f6f2a9579f 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -149,8 +149,6 @@ async def async_firewall_policy_control_fn( policy = hub.api.firewall_policies[obj_id].raw policy["enabled"] = target await hub.api.request(FirewallPolicyUpdateRequest.create(policy)) - # Update the policies so the UI is updated appropriately - await hub.api.firewall_policies.update() @callback @@ -460,10 +458,7 @@ class UnifiSwitchEntity[HandlerT: APIHandler, ApiItemT: ApiItem]( translation_domain=DOMAIN, translation_key="action_request_failed", ) from err - if coordinator := self.hub.entity_loader.get_data_update_coordinator( - self.entity_description.api_handler_fn(self.api) - ): - await coordinator.async_request_refresh() + await self.async_refresh_after_control() @override async def async_turn_off(self, **kwargs: Any) -> None: @@ -475,10 +470,7 @@ class UnifiSwitchEntity[HandlerT: APIHandler, ApiItemT: ApiItem]( translation_domain=DOMAIN, translation_key="action_request_failed", ) from err - if coordinator := self.hub.entity_loader.get_data_update_coordinator( - self.entity_description.api_handler_fn(self.api) - ): - await coordinator.async_request_refresh() + await self.async_refresh_after_control() @callback @override diff --git a/tests/components/unifi/test_hub.py b/tests/components/unifi/test_hub.py index 78f9d484619c..ab3821fdb851 100644 --- a/tests/components/unifi/test_hub.py +++ b/tests/components/unifi/test_hub.py @@ -6,20 +6,27 @@ from typing import Any from unittest.mock import patch import aiounifi +from aiounifi.interfaces.api_handlers import ItemEvent +from aiounifi.models.message import MessageKey import pytest -from homeassistant.components.unifi.const import DOMAIN +from homeassistant.components.unifi.const import CONF_BLOCK_CLIENT, DOMAIN +from homeassistant.components.unifi.coordinator import POLL_INTERVAL from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect from homeassistant.components.unifi.hub import get_unifi_api from homeassistant.config_entries import ConfigEntryState -from homeassistant.const import CONF_HOST, Platform -from homeassistant.core import HomeAssistant +from homeassistant.const import CONF_HOST, EVENT_STATE_REPORTED, Platform +from homeassistant.core import Event, EventStateReportedData, HomeAssistant, callback from homeassistant.helpers import device_registry as dr from homeassistant.util import dt as dt_util -from .conftest import ConfigEntryFactoryType, WebsocketStateManager +from .conftest import ( + ConfigEntryFactoryType, + WebsocketMessageMock, + WebsocketStateManager, +) -from tests.common import MockConfigEntry +from tests.common import MockConfigEntry, async_fire_time_changed from tests.test_util.aiohttp import AiohttpClientMocker @@ -56,6 +63,305 @@ async def test_hub_setup( assert device_entry.sw_version == "7.4.162" +async def test_coordinators_preserve_handler_update_sources( + config_entry_setup: MockConfigEntry, +) -> None: + """Ensure coordinator polling matches the handler's existing update source.""" + loader = config_entry_setup.runtime_data.entity_loader + api = config_entry_setup.runtime_data.api + + clients_coordinator = loader.get_data_update_coordinator(api.clients) + devices_coordinator = loader.get_data_update_coordinator(api.devices) + assert clients_coordinator.update_interval is None + assert devices_coordinator.update_interval is None + + assert loader.get_data_update_coordinator(api.ports) is devices_coordinator + assert loader.get_data_update_coordinator(api.outlets) is devices_coordinator + + for handler in ( + api.object_oriented_network_configs, + api.traffic_rules, + api.traffic_routes, + ): + coordinator = loader.get_data_update_coordinator(handler) + assert coordinator.update_interval == POLL_INTERVAL + + +async def test_get_data_update_coordinator_requires_registered_handler( + config_entry_setup: MockConfigEntry, +) -> None: + """Ensure a handler without a coordinator fails at lookup time.""" + loader = config_entry_setup.runtime_data.entity_loader + api = config_entry_setup.runtime_data.api + + with pytest.raises(KeyError): + loader.get_data_update_coordinator(api.sites) + + +async def test_polling_coordinator_refreshes_after_interval( + hass: HomeAssistant, + config_entry_setup: MockConfigEntry, +) -> None: + """Ensure polling coordinators refresh when their interval elapses.""" + loader = config_entry_setup.runtime_data.entity_loader + api = config_entry_setup.runtime_data.api + coordinator = loader.get_data_update_coordinator( + api.object_oriented_network_configs + ) + + assert coordinator.update_interval == POLL_INTERVAL + + with patch.object( + api.object_oriented_network_configs, + "update", + wraps=api.object_oriented_network_configs.update, + ) as mock_update: + async_fire_time_changed(hass, dt_util.utcnow() + POLL_INTERVAL) + await hass.async_block_till_done() + + assert mock_update.call_count >= 1 + + +@pytest.mark.parametrize( + "object_oriented_network_config_payload", + [ + [ + { + "id": "69f6b0a5e0e3ee2d4614cb5c", + "enabled": True, + "name": "Nintendo Switch - Block Internet", + "target_type": "CLIENTS", + "targets": ["00:00:00:00:00:01"], + "qos": {"enabled": False}, + "route": {"enabled": False}, + "secure": { + "enabled": True, + "internet": { + "mode": "TURN_OFF_INTERNET", + "schedule": {"mode": "ALWAYS"}, + }, + }, + } + ] + ], +) +async def test_entity_unavailable_on_polling_coordinator_failure( + hass: HomeAssistant, + config_entry_setup: MockConfigEntry, +) -> None: + """Ensure polling failures make entities unavailable until recovery.""" + entity_id = "switch.unifi_network_nintendo_switch_block_internet" + state = hass.states.get(entity_id) + assert state is not None + assert state.state != "unavailable" + + coordinator = ( + config_entry_setup.runtime_data.entity_loader.get_data_update_coordinator( + config_entry_setup.runtime_data.api.object_oriented_network_configs + ) + ) + + with patch.object( + coordinator.handler, + "update", + side_effect=RuntimeError("Polling error"), + ): + await coordinator.async_refresh() + await hass.async_block_till_done() + + assert coordinator.last_update_success is False + state = hass.states.get(entity_id) + assert state is not None + assert state.state == "unavailable" + + await coordinator.async_refresh() + await hass.async_block_till_done() + + assert coordinator.last_update_success is True + state = hass.states.get(entity_id) + assert state is not None + assert state.state != "unavailable" + + +async def test_websocket_updates_notify_coordinator( + config_entry_setup: MockConfigEntry, + mock_websocket_message: WebsocketMessageMock, +) -> None: + """Ensure websocket handler updates are forwarded through the coordinator.""" + coordinator = ( + config_entry_setup.runtime_data.entity_loader.get_data_update_coordinator( + config_entry_setup.runtime_data.api.clients + ) + ) + assert coordinator is not None + + with patch.object( + coordinator, "async_set_updated_data", wraps=coordinator.async_set_updated_data + ) as set_updated_data: + mock_websocket_message( + message=MessageKey.CLIENT, + data={ + "hostname": "client", + "ip": "10.0.0.1", + "is_wired": True, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:01", + "name": "Client", + }, + ) + + set_updated_data.assert_called_once_with((ItemEvent.ADDED, "00:00:00:00:00:01")) + + +async def test_coordinator_filters_websocket_listeners_and_broadcasts_polling( + hass: HomeAssistant, + config_entry_setup: MockConfigEntry, + mock_websocket_message: WebsocketMessageMock, +) -> None: + """Ensure websocket updates filter listeners and polling updates broadcast.""" + clients_coordinator = ( + config_entry_setup.runtime_data.entity_loader.get_data_update_coordinator( + config_entry_setup.runtime_data.api.clients + ) + ) + client_calls: list[str] = [] + other_calls: list[str] = [] + + @callback + def client_listener() -> None: + client_calls.append("called") + + @callback + def other_listener() -> None: + other_calls.append("called") + + remove_client_listener = clients_coordinator.async_add_listener( + client_listener, context=("00:00:00:00:00:01",) + ) + remove_other_listener = clients_coordinator.async_add_listener( + other_listener, context=("00:00:00:00:00:02",) + ) + + mock_websocket_message( + message=MessageKey.CLIENT, + data={ + "hostname": "client", + "ip": "10.0.0.1", + "is_wired": True, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:01", + "name": "Client", + }, + ) + + assert client_calls == ["called"] + assert other_calls == [] + + polling_coordinator = ( + config_entry_setup.runtime_data.entity_loader.get_data_update_coordinator( + config_entry_setup.runtime_data.api.object_oriented_network_configs + ) + ) + polling_calls: list[str] = [] + other_polling_calls: list[str] = [] + + @callback + def polling_listener() -> None: + polling_calls.append("called") + + @callback + def other_polling_listener() -> None: + other_polling_calls.append("called") + + remove_polling_listener = polling_coordinator.async_add_listener( + polling_listener, context=("any-object",) + ) + remove_other_polling_listener = polling_coordinator.async_add_listener( + other_polling_listener, context=("another-object",) + ) + polling_coordinator.async_set_updated_data(None) + await hass.async_block_till_done() + + assert polling_calls == ["called"] + assert other_polling_calls == ["called"] + + remove_client_listener() + remove_other_listener() + remove_polling_listener() + remove_other_polling_listener() + + +@pytest.mark.parametrize( + "config_entry_options", + [{CONF_BLOCK_CLIENT: ["00:00:00:00:00:01", "00:00:00:00:00:02"]}], +) +@pytest.mark.parametrize( + "client_payload", + [ + [ + { + "blocked": True, + "hostname": "client_1", + "ip": "10.0.0.1", + "is_wired": True, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:01", + "name": "Client 1", + }, + { + "blocked": True, + "hostname": "client_2", + "ip": "10.0.0.2", + "is_wired": True, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:02", + "name": "Client 2", + }, + ] + ], +) +async def test_coordinator_update_only_refreshes_changed_entity( + hass: HomeAssistant, + config_entry_setup: MockConfigEntry, + mock_websocket_message: WebsocketMessageMock, +) -> None: + """Ensure a coordinator update for one object does not refresh unrelated entities.""" + changed_entity_id = "switch.client_1_blocked" + other_entity_id = "switch.client_2_blocked" + assert hass.states.get(changed_entity_id) is not None + assert hass.states.get(other_entity_id) is not None + + written_entity_ids: list[str] = [] + + @callback + def track_state_reported(event: Event[EventStateReportedData]) -> None: + written_entity_ids.append(event.data["entity_id"]) + + @callback + def filter_tracked_entities(data: EventStateReportedData) -> bool: + return data["entity_id"] in (changed_entity_id, other_entity_id) + + hass.bus.async_listen( + EVENT_STATE_REPORTED, track_state_reported, filter_tracked_entities + ) + + mock_websocket_message( + message=MessageKey.CLIENT, + data={ + "hostname": "client_1", + "ip": "10.0.0.1", + "is_wired": True, + "last_seen": 1562600146, + "mac": "00:00:00:00:00:01", + "name": "Client 1", + }, + ) + await hass.async_block_till_done() + + assert changed_entity_id in written_entity_ids + assert other_entity_id not in written_entity_ids + + async def test_reset_after_successful_setup( hass: HomeAssistant, config_entry_setup: MockConfigEntry ) -> None: diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 4b223fde41f0..fbf458776574 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -1292,7 +1292,7 @@ async def test_traffic_rules( expected_enable_call = deepcopy(traffic_rule) expected_enable_call["enabled"] = True - assert aioclient_mock.call_count == call_count + 1 + assert aioclient_mock.call_count == call_count + 2 assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call @@ -1347,7 +1347,7 @@ async def test_traffic_routes( expected_enable_call = deepcopy(traffic_route) expected_enable_call["enabled"] = True - assert aioclient_mock.call_count == call_count + 1 + assert aioclient_mock.call_count == call_count + 2 assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call @@ -1436,13 +1436,25 @@ async def test_object_oriented_network_configs( aioclient_mock.put(config_url) call_count = aioclient_mock.call_count - - await hass.services.async_call( - SWITCH_DOMAIN, - "turn_off", - {"entity_id": entity_id}, - blocking=True, + coordinator = ( + config_entry_setup.runtime_data.entity_loader.get_data_update_coordinator( + config_entry_setup.runtime_data.api.object_oriented_network_configs + ) ) + + with ( + patch.object(coordinator, "async_refresh") as async_refresh, + patch.object(coordinator, "async_request_refresh") as async_request_refresh, + ): + await hass.services.async_call( + SWITCH_DOMAIN, + "turn_off", + {"entity_id": entity_id}, + blocking=True, + ) + + async_refresh.assert_awaited_once() + async_request_refresh.assert_not_awaited() expected_disable_call = deepcopy(config) expected_disable_call["enabled"] = False