Rework UniFi wireless client "wired bug" logic (#89757)

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]:
+    def is_wireless(self, client: Client) -> bool:
-        """Get data related to a specific controller."""
+        """Is client known to be wireless.
-        data = self.data.get(config_entry.entry_id, {"wireless_devices": []})
+
-        return set(data["wireless_devices"])
+        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

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 = hass.data[UNIFI_WIRELESS_CLIENTS]
-        self.wireless_clients = None
 
         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.update_clients(set(self.api.clients.values()))
-        self.wireless_clients = wireless_clients.get_data(self.config_entry)
-        self.update_wireless_clients()
 
         self.config_entry.add_update_listener(self.async_config_entry_updated)
 

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
 

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:
+    if controller.wireless_clients.is_wireless(client):
-        return client.wired_rx_bytes_r / 1000000
+        return client.rx_bytes_r / 1000000
-    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:
+    if controller.wireless_clients.is_wireless(client):
-        return client.wired_tx_bytes_r / 1000000
+        return client.tx_bytes_r / 1000000
-    return client.tx_bytes_r / 1000000
+    return client.wired_tx_bytes_r / 1000000
 
 
 @callback

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(
+        hass.data[UNIFI_WIRELESS_CLIENTS].wireless_clients.update(
-            known_wireless_clients, config_entry
+            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:

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"
-            ]
-        )