Move wireless clients into entity_loader (#112813)

homeassistant/components/unifi/device_tracker.py

@@ -89,7 +89,7 @@ def async_client_allowed_fn(hub: UnifiHub, obj_id: str) -> bool:
         return False
 
     client = hub.api.clients[obj_id]
-    if client.mac not in hub.wireless_clients:
+    if client.mac not in hub.entity_loader.wireless_clients:
         if not hub.config.option_track_wired_clients:
             return False
 
@@ -108,7 +108,7 @@ def async_client_is_connected_fn(hub: UnifiHub, obj_id: str) -> bool:
     """Check if device object is disabled."""
     client = hub.api.clients[obj_id]
 
-    if hub.wireless_clients.is_wireless(client) and client.is_wired:
+    if hub.entity_loader.wireless_clients.is_wireless(client) and client.is_wired:
         if not hub.config.option_ignore_wired_bug:
             return False  # Wired bug in action
 

homeassistant/components/unifi/hub/entity_loader.py

@@ -20,7 +20,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import async_entries_for_config_entry
 
-from ..const import LOGGER
+from ..const import LOGGER, UNIFI_WIRELESS_CLIENTS
 from ..entity import UnifiEntity, UnifiEntityDescription
 
 if TYPE_CHECKING:
@@ -46,6 +46,7 @@ class UnifiEntityLoader:
             hub.api.system_information.update,
             hub.api.wlans.update,
         )
+        self.wireless_clients = hub.hass.data[UNIFI_WIRELESS_CLIENTS]
 
         self.platforms: list[
             tuple[
@@ -63,6 +64,7 @@ class UnifiEntityLoader:
         """Initialize API data and extra client support."""
         await self.refresh_api_data()
         self.restore_inactive_clients()
+        self.wireless_clients.update_clients(set(self.hub.api.clients.values()))
 
     async def refresh_api_data(self) -> None:
         """Refresh API data from network application."""

homeassistant/components/unifi/hub/hub.py

@@ -19,13 +19,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.event import async_call_later, async_track_time_interval
 import homeassistant.util.dt as dt_util
 
-from ..const import (
-    ATTR_MANUFACTURER,
-    CONF_SITE_ID,
-    DOMAIN as UNIFI_DOMAIN,
-    PLATFORMS,
-    UNIFI_WIRELESS_CLIENTS,
-)
+from ..const import ATTR_MANUFACTURER, CONF_SITE_ID, DOMAIN as UNIFI_DOMAIN, PLATFORMS
 from .config import UnifiConfig
 from .entity_loader import UnifiEntityLoader
 from .websocket import UnifiWebsocket
@@ -46,8 +40,6 @@ class UnifiHub:
         self.entity_loader = UnifiEntityLoader(self)
         self.websocket = UnifiWebsocket(hass, api, self.signal_reachable)
 
-        self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
-
         self.site = config_entry.data[CONF_SITE_ID]
         self.is_admin = False
 
@@ -91,8 +83,6 @@ class UnifiHub:
         assert self.config.entry.unique_id is not None
         self.is_admin = self.api.sites[self.config.entry.unique_id].role == "admin"
 
-        self.wireless_clients.update_clients(set(self.api.clients.values()))
-
         self.config.entry.add_update_listener(self.async_config_entry_updated)
 
         self._cancel_heartbeat_check = async_track_time_interval(

homeassistant/components/unifi/sensor.py

@@ -72,7 +72,7 @@ def async_uptime_sensor_allowed_fn(hub: UnifiHub, obj_id: str) -> bool:
 @callback
 def async_client_rx_value_fn(hub: UnifiHub, client: Client) -> float:
     """Calculate receiving data transfer value."""
-    if hub.wireless_clients.is_wireless(client):
+    if hub.entity_loader.wireless_clients.is_wireless(client):
         return client.rx_bytes_r / 1000000
     return client.wired_rx_bytes_r / 1000000
 
@@ -80,7 +80,7 @@ def async_client_rx_value_fn(hub: UnifiHub, client: Client) -> float:
 @callback
 def async_client_tx_value_fn(hub: UnifiHub, client: Client) -> float:
     """Calculate transmission data transfer value."""
-    if hub.wireless_clients.is_wireless(client):
+    if hub.entity_loader.wireless_clients.is_wireless(client):
         return client.tx_bytes_r / 1000000
     return client.wired_tx_bytes_r / 1000000