diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 8aae95bda41..7e46ff2d7e0 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -9,6 +9,7 @@ from typing import Any from aiohttp import CookieJar import aiounifi +from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.messages import DATA_CLIENT_REMOVED, DATA_EVENT from aiounifi.models.event import EventKey from aiounifi.websocket import WebsocketSignal, WebsocketState @@ -31,6 +32,7 @@ from homeassistant.helpers import ( ) from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from homeassistant.helpers.dispatcher import async_dispatcher_send +from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_registry import async_entries_for_config_entry from homeassistant.helpers.event import async_track_time_interval import homeassistant.util.dt as dt_util @@ -62,6 +64,7 @@ from .const import ( PLATFORMS, UNIFI_WIRELESS_CLIENTS, ) +from .entity import UnifiEntity, UnifiEntityDescription from .errors import AuthenticationRequired, CannotConnect RETRY_TIMER = 15 @@ -183,6 +186,44 @@ class UniFiController: return client.mac return None + @callback + def register_platform_add_entities( + self, + unifi_platform_entity: type[UnifiEntity], + descriptions: tuple[UnifiEntityDescription, ...], + async_add_entities: AddEntitiesCallback, + ) -> None: + """Subscribe to UniFi API handlers and create entities.""" + + @callback + def async_load_entities(description: UnifiEntityDescription) -> None: + """Load and subscribe to UniFi endpoints.""" + entities: list[UnifiEntity] = [] + api_handler = description.api_handler_fn(self.api) + + @callback + def async_create_entity(event: ItemEvent, obj_id: str) -> None: + """Create UniFi entity.""" + if not description.allowed_fn( + self, obj_id + ) or not description.supported_fn(self.api, obj_id): + return + + entity = unifi_platform_entity(obj_id, self, description) + if event == ItemEvent.ADDED: + async_add_entities([entity]) + return + entities.append(entity) + + for obj_id in api_handler: + async_create_entity(ItemEvent.CHANGED, obj_id) + async_add_entities(entities) + + api_handler.subscribe(async_create_entity, ItemEvent.ADDED) + + for description in descriptions: + async_load_entities(description) + @callback def async_unifi_signalling_callback(self, signal, data): """Handle messages back from UniFi library.""" diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index 79a80fad73c..b7aed362133 100644 --- a/homeassistant/components/unifi/entity.py +++ b/homeassistant/components/unifi/entity.py @@ -4,7 +4,7 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import Callable from dataclasses import dataclass -from typing import Generic, TypeVar +from typing import TYPE_CHECKING, Generic, TypeVar import aiounifi from aiounifi.interfaces.api_handlers import CallbackType, ItemEvent, UnsubscribeType @@ -17,7 +17,8 @@ from homeassistant.helpers import entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import DeviceInfo, Entity, EntityDescription -from .controller import UniFiController +if TYPE_CHECKING: + from .controller import UniFiController DataT = TypeVar("DataT", bound=Device) HandlerT = TypeVar("HandlerT", bound=Devices) diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index 6cff6b7932d..36a94b73da2 100644 --- a/homeassistant/components/unifi/update.py +++ b/homeassistant/components/unifi/update.py @@ -103,35 +103,9 @@ async def async_setup_entry( ) -> None: """Set up update entities for UniFi Network integration.""" controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] - - @callback - def async_load_entities(description: UnifiUpdateEntityDescription) -> None: - """Load and subscribe to UniFi devices.""" - entities: list[UpdateEntity] = [] - api_handler = description.api_handler_fn(controller.api) - - @callback - def async_create_entity(event: ItemEvent, obj_id: str) -> None: - """Create UniFi entity.""" - if not description.allowed_fn( - controller, obj_id - ) or not description.supported_fn(controller.api, obj_id): - return - - entity = UnifiDeviceUpdateEntity(obj_id, controller, description) - if event == ItemEvent.ADDED: - async_add_entities([entity]) - return - entities.append(entity) - - for obj_id in api_handler: - async_create_entity(ItemEvent.CHANGED, obj_id) - async_add_entities(entities) - - api_handler.subscribe(async_create_entity, ItemEvent.ADDED) - - for description in ENTITY_DESCRIPTIONS: - async_load_entities(description) + controller.register_platform_add_entities( + UnifiDeviceUpdateEntity, ENTITY_DESCRIPTIONS, async_add_entities + ) class UnifiDeviceUpdateEntity(UnifiEntity[HandlerT, DataT], UpdateEntity):