diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index adaa7c977f7..d6405d11716 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -1,5 +1,7 @@ """Integration to UniFi Network and its various features.""" +from aiounifi.models.client import Client + from homeassistant.config_entries import ConfigEntry from homeassistant.const import EVENT_HOMEASSISTANT_STOP from homeassistant.core import HomeAssistant, callback @@ -91,33 +93,55 @@ def async_remove_poe_client_entities( class UnifiWirelessClients: """Class to store clients known to be wireless. - This is needed since wireless devices going offline might get marked as wired by UniFi. + This is needed since wireless devices going offline + might get marked as wired by UniFi. """ def __init__(self, hass: HomeAssistant) -> None: """Set up client storage.""" self.hass = hass - self.data: dict[str, dict[str, list[str]]] = {} + self.data: dict[str, dict[str, list[str]] | list[str]] = {} + self.wireless_clients: set[str] = set() self._store: Store = Store(hass, STORAGE_VERSION, STORAGE_KEY) async def async_load(self) -> None: """Load data from file.""" if (data := await self._store.async_load()) is not None: self.data = data + if "wireless_clients" not in data: + data["wireless_clients"] = [ + obj_id + for config_entry in data + for obj_id in data[config_entry]["wireless_devices"] + ] + self.wireless_clients.update(data["wireless_clients"]) @callback - def get_data(self, config_entry: ConfigEntry) -> set[str]: - """Get data related to a specific controller.""" - data = self.data.get(config_entry.entry_id, {"wireless_devices": []}) - return set(data["wireless_devices"]) + def is_wireless(self, client: Client) -> bool: + """Is client known to be wireless. + + Store if client is wireless and not known. + """ + if not client.is_wired and client.mac not in self.wireless_clients: + self.wireless_clients.add(client.mac) + self._store.async_delay_save(self._data_to_save, SAVE_DELAY) + + return client.mac in self.wireless_clients @callback - def update_data(self, data: set[str], config_entry: ConfigEntry) -> None: + def update_clients(self, clients: set[Client]) -> None: """Update data and schedule to save to file.""" - self.data[config_entry.entry_id] = {"wireless_devices": list(data)} + self.wireless_clients.update( + {client.mac for client in clients if not client.is_wired} + ) self._store.async_delay_save(self._data_to_save, SAVE_DELAY) @callback - def _data_to_save(self) -> dict[str, dict[str, list[str]]]: + def _data_to_save(self) -> dict[str, dict[str, list[str]] | list[str]]: """Return data of UniFi wireless clients to store in a file.""" + self.data["wireless_clients"] = list(self.wireless_clients) return self.data + + def __contains__(self, obj_id: int | str) -> bool: + """Validate membership of item ID.""" + return obj_id in self.wireless_clients diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 8a047606c67..a5f3c4d7720 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -10,8 +10,6 @@ from typing import Any from aiohttp import CookieJar import aiounifi from aiounifi.interfaces.api_handlers import ItemEvent -from aiounifi.interfaces.messages import DATA_EVENT -from aiounifi.models.event import EventKey from aiounifi.websocket import WebsocketSignal, WebsocketState import async_timeout @@ -86,8 +84,7 @@ class UniFiController: api.callback = self.async_unifi_signalling_callback self.available = True - self.progress = None - self.wireless_clients = None + self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS] self.site_id: str = "" self._site_name = None @@ -247,15 +244,6 @@ class UniFiController: else: LOGGER.info("Connected to UniFi Network") - elif signal == WebsocketSignal.DATA and DATA_EVENT in data: - for event in data[DATA_EVENT]: - if event.key in ( - EventKey.WIRELESS_CLIENT_CONNECTED, - EventKey.WIRELESS_GUEST_CONNECTED, - ): - self.update_wireless_clients() - break - @property def signal_reachable(self) -> str: """Integration specific event to signal a change in connection status.""" @@ -271,22 +259,6 @@ class UniFiController: """Event specific per UniFi device tracker to signal new heartbeat missed.""" return "unifi-heartbeat-missed" - def update_wireless_clients(self): - """Update set of known to be wireless clients.""" - new_wireless_clients = set() - - for client_id in self.api.clients: - if ( - client_id not in self.wireless_clients - and not self.api.clients[client_id].is_wired - ): - new_wireless_clients.add(client_id) - - if new_wireless_clients: - self.wireless_clients |= new_wireless_clients - unifi_wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS] - unifi_wireless_clients.update_data(self.wireless_clients, self.config_entry) - async def initialize(self): """Set up a UniFi Network instance.""" await self.api.initialize() @@ -326,9 +298,7 @@ class UniFiController: client.mac, ) - wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS] - self.wireless_clients = wireless_clients.get_data(self.config_entry) - self.update_wireless_clients() + self.wireless_clients.update_clients(set(self.api.clients.values())) self.config_entry.add_update_listener(self.async_config_entry_updated) diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index a5b153d7f36..f891416c635 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -105,7 +105,7 @@ def async_client_is_connected_fn(controller: UniFiController, obj_id: str) -> bo """Check if device object is disabled.""" client = controller.api.clients[obj_id] - if client.is_wired != (obj_id not in controller.wireless_clients): + if controller.wireless_clients.is_wireless(client) and client.is_wired: if not controller.option_ignore_wired_bug: return False # Wired bug in action diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 05598589feb..420fc3803c3 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -45,17 +45,17 @@ from .entity import ( @callback def async_client_rx_value_fn(controller: UniFiController, client: Client) -> float: """Calculate receiving data transfer value.""" - if client.mac not in controller.wireless_clients: - return client.wired_rx_bytes_r / 1000000 - return client.rx_bytes_r / 1000000 + if controller.wireless_clients.is_wireless(client): + return client.rx_bytes_r / 1000000 + return client.wired_rx_bytes_r / 1000000 @callback def async_client_tx_value_fn(controller: UniFiController, client: Client) -> float: """Calculate transmission data transfer value.""" - if client.mac not in controller.wireless_clients: - return client.wired_tx_bytes_r / 1000000 - return client.tx_bytes_r / 1000000 + if controller.wireless_clients.is_wireless(client): + return client.tx_bytes_r / 1000000 + return client.wired_tx_bytes_r / 1000000 @callback diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 931c0fccdf0..e3efaef915b 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -6,8 +6,6 @@ from http import HTTPStatus from unittest.mock import Mock, patch import aiounifi -from aiounifi.models.event import EventKey -from aiounifi.models.message import MessageKey from aiounifi.websocket import WebsocketState import pytest @@ -182,8 +180,8 @@ async def setup_unifi_integration( config_entry.add_to_hass(hass) if known_wireless_clients: - hass.data[UNIFI_WIRELESS_CLIENTS].update_data( - known_wireless_clients, config_entry + hass.data[UNIFI_WIRELESS_CLIENTS].wireless_clients.update( + known_wireless_clients ) if aioclient_mock: @@ -383,41 +381,6 @@ async def test_connection_state_signalling( assert hass.states.get("device_tracker.client").state == "home" -async def test_wireless_client_event_calls_update_wireless_devices( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket -) -> None: - """Call update_wireless_devices method when receiving wireless client event.""" - client_1_dict = { - "essid": "ssid", - "disabled": False, - "hostname": "client_1", - "ip": "10.0.0.4", - "is_wired": False, - "last_seen": dt_util.as_timestamp(dt_util.utcnow()), - "mac": "00:00:00:00:00:01", - } - await setup_unifi_integration( - hass, - aioclient_mock, - clients_response=[client_1_dict], - known_wireless_clients=(client_1_dict["mac"],), - ) - - with patch( - "homeassistant.components.unifi.controller.UniFiController.update_wireless_clients", - return_value=None, - ) as wireless_clients_mock: - event = { - "datetime": "2020-01-20T19:37:04Z", - "user": "00:00:00:00:00:01", - "key": EventKey.WIRELESS_CLIENT_CONNECTED.value, - "msg": "User[11:22:33:44:55:66] has connected to WLAN", - "time": 1579549024893, - } - mock_unifi_websocket(message=MessageKey.EVENT, data=event) - assert wireless_clients_mock.assert_called_once - - async def test_reconnect_mechanism( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket ) -> None: diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index cb232445eb1..cce26ac84cc 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -89,19 +89,13 @@ async def test_wireless_clients( "is_wired": False, "mac": "00:00:00:00:00:02", } - config_entry = await setup_unifi_integration( + await setup_unifi_integration( hass, aioclient_mock, clients_response=[client_1, client_2] ) await flush_store(hass.data[unifi.UNIFI_WIRELESS_CLIENTS]._store) - for mac in [ + assert sorted(hass_storage[unifi.STORAGE_KEY]["data"]["wireless_clients"]) == [ "00:00:00:00:00:00", "00:00:00:00:00:01", "00:00:00:00:00:02", - ]: - assert ( - mac - in hass_storage[unifi.STORAGE_KEY]["data"][config_entry.entry_id][ - "wireless_devices" - ] - ) + ]