From f2879e6f393a6e655003ece9d0f7714d2b4d52b0 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Sat, 9 Mar 2024 09:18:23 +0100 Subject: [PATCH] Break out UniFi platform registration to its own class (#112514) --- homeassistant/components/unifi/__init__.py | 1 + homeassistant/components/unifi/button.py | 4 +- .../components/unifi/device_tracker.py | 4 +- homeassistant/components/unifi/entity.py | 6 +- .../components/unifi/hub/entity_loader.py | 131 ++++++++++++++++++ homeassistant/components/unifi/hub/hub.py | 101 ++------------ homeassistant/components/unifi/image.py | 4 +- homeassistant/components/unifi/sensor.py | 4 +- homeassistant/components/unifi/switch.py | 4 +- homeassistant/components/unifi/update.py | 4 +- tests/components/unifi/test_config_flow.py | 12 +- 11 files changed, 158 insertions(+), 117 deletions(-) create mode 100644 homeassistant/components/unifi/hub/entity_loader.py diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index dda91801084..5174a1a7796 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -49,6 +49,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS) hub.async_update_device_registry() + hub.entity_loader.load_entities() if len(hass.data[UNIFI_DOMAIN]) == 1: async_setup_services(hass) diff --git a/homeassistant/components/unifi/button.py b/homeassistant/components/unifi/button.py index 950ad2f8361..c8f5d5b33f7 100644 --- a/homeassistant/components/unifi/button.py +++ b/homeassistant/components/unifi/button.py @@ -107,9 +107,7 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up button platform for UniFi Network integration.""" - UnifiHub.register_platform( - hass, - config_entry, + UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform( async_add_entities, UnifiButtonEntity, ENTITY_DESCRIPTIONS, diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index cce41135e4c..d84691b814c 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -220,8 +220,8 @@ async def async_setup_entry( ) -> None: """Set up device tracker for UniFi Network integration.""" async_update_unique_id(hass, config_entry) - UnifiHub.register_platform( - hass, config_entry, async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS + UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform( + async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS ) diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index 5f63fdb1980..5a626ec56f6 100644 --- a/homeassistant/components/unifi/entity.py +++ b/homeassistant/components/unifi/entity.py @@ -133,7 +133,7 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]): self.hub = hub self.entity_description = description - hub.known_objects.add((description.key, obj_id)) + hub.entity_loader.known_objects.add((description.key, obj_id)) self._removed = False @@ -154,7 +154,9 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]): @callback def unregister_object() -> None: """Remove object ID from known_objects when unloaded.""" - self.hub.known_objects.discard((description.key, self._obj_id)) + self.hub.entity_loader.known_objects.discard( + (description.key, self._obj_id) + ) self.async_on_remove(unregister_object) diff --git a/homeassistant/components/unifi/hub/entity_loader.py b/homeassistant/components/unifi/hub/entity_loader.py new file mode 100644 index 00000000000..be053403bff --- /dev/null +++ b/homeassistant/components/unifi/hub/entity_loader.py @@ -0,0 +1,131 @@ +"""UniFi Network entity loader. + +Central point to load entities for the different platforms. +Make sure expected clients are available for platforms. +""" +from __future__ import annotations + +from collections.abc import Iterable +from datetime import timedelta +from functools import partial +from typing import TYPE_CHECKING + +from aiounifi.interfaces.api_handlers import ItemEvent + +from homeassistant.core import callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_platform import AddEntitiesCallback + +from ..entity import UnifiEntity, UnifiEntityDescription + +if TYPE_CHECKING: + from .hub import UnifiHub + +CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1) + + +class UnifiEntityLoader: + """UniFi Network integration handling platforms for entity registration.""" + + def __init__(self, hub: UnifiHub) -> None: + """Initialize the UniFi entity loader.""" + self.hub = hub + + self.platforms: list[ + tuple[ + AddEntitiesCallback, + type[UnifiEntity], + tuple[UnifiEntityDescription, ...], + bool, + ] + ] = [] + + self.known_objects: set[tuple[str, str]] = set() + """Tuples of entity description key and object ID of loaded entities.""" + + @callback + def register_platform( + self, + async_add_entities: AddEntitiesCallback, + entity_class: type[UnifiEntity], + descriptions: tuple[UnifiEntityDescription, ...], + requires_admin: bool = False, + ) -> None: + """Register UniFi entity platforms.""" + self.platforms.append( + (async_add_entities, entity_class, descriptions, requires_admin) + ) + + @callback + def load_entities(self) -> None: + """Populate UniFi platforms with entities.""" + for ( + async_add_entities, + entity_class, + descriptions, + requires_admin, + ) in self.platforms: + if requires_admin and not self.hub.is_admin: + continue + self._load_entities(entity_class, descriptions, async_add_entities) + + @callback + def _should_add_entity( + self, description: UnifiEntityDescription, obj_id: str + ) -> bool: + """Check if entity should be added.""" + return bool( + (description.key, obj_id) not in self.known_objects + and description.allowed_fn(self.hub, obj_id) + and description.supported_fn(self.hub, obj_id) + ) + + @callback + def _load_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(descriptions: Iterable[UnifiEntityDescription]) -> None: + """Load and subscribe to UniFi endpoints.""" + + @callback + def _add_unifi_entities() -> None: + """Add UniFi entity.""" + async_add_entities( + unifi_platform_entity(obj_id, self.hub, description) + for description in descriptions + for obj_id in description.api_handler_fn(self.hub.api) + if self._should_add_entity(description, obj_id) + ) + + _add_unifi_entities() + + @callback + def _create_unifi_entity( + description: UnifiEntityDescription, event: ItemEvent, obj_id: str + ) -> None: + """Create new UniFi entity on event.""" + if self._should_add_entity(description, obj_id): + async_add_entities( + [unifi_platform_entity(obj_id, self.hub, description)] + ) + + for description in descriptions: + description.api_handler_fn(self.hub.api).subscribe( + partial(_create_unifi_entity, description), ItemEvent.ADDED + ) + + self.hub.config.entry.async_on_unload( + async_dispatcher_connect( + self.hub.hass, + self.hub.signal_options_update, + _add_unifi_entities, + ) + ) + + async_load_entities(descriptions) diff --git a/homeassistant/components/unifi/hub/hub.py b/homeassistant/components/unifi/hub/hub.py index 89e741c43d5..ae485ec36b8 100644 --- a/homeassistant/components/unifi/hub/hub.py +++ b/homeassistant/components/unifi/hub/hub.py @@ -2,12 +2,9 @@ from __future__ import annotations -from collections.abc import Iterable from datetime import datetime, timedelta -from functools import partial import aiounifi -from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.models.device import DeviceSetPoePortModeRequest from homeassistant.config_entries import ConfigEntry @@ -19,11 +16,7 @@ from homeassistant.helpers.device_registry import ( DeviceEntryType, DeviceInfo, ) -from homeassistant.helpers.dispatcher import ( - async_dispatcher_connect, - async_dispatcher_send, -) -from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.entity_registry import async_entries_for_config_entry from homeassistant.helpers.event import async_call_later, async_track_time_interval import homeassistant.util.dt as dt_util @@ -35,8 +28,8 @@ from ..const import ( PLATFORMS, UNIFI_WIRELESS_CLIENTS, ) -from ..entity import UnifiEntity, UnifiEntityDescription from .config import UnifiConfig +from .entity_loader import UnifiEntityLoader from .websocket import UnifiWebsocket CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1) @@ -52,6 +45,7 @@ class UnifiHub: self.hass = hass self.api = api self.config = UnifiConfig.from_config_entry(config_entry) + self.entity_loader = UnifiEntityLoader(self) self.websocket = UnifiWebsocket(hass, api, self.signal_reachable) self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS] @@ -62,96 +56,21 @@ class UnifiHub: self._cancel_heartbeat_check: CALLBACK_TYPE | None = None self._heartbeat_time: dict[str, datetime] = {} - self.entities: dict[str, str] = {} - self.known_objects: set[tuple[str, str]] = set() - self.poe_command_queue: dict[str, dict[int, str]] = {} self._cancel_poe_command: CALLBACK_TYPE | None = None + @callback + @staticmethod + def get_hub(hass: HomeAssistant, config_entry: ConfigEntry) -> UnifiHub: + """Get UniFi hub from config entry.""" + hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + return hub + @property def available(self) -> bool: """Websocket connection state.""" return self.websocket.available - @callback - @staticmethod - def register_platform( - hass: HomeAssistant, - config_entry: ConfigEntry, - async_add_entities: AddEntitiesCallback, - entity_class: type[UnifiEntity], - descriptions: tuple[UnifiEntityDescription, ...], - requires_admin: bool = False, - ) -> None: - """Register platform for UniFi entity management.""" - hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id] - if requires_admin and not hub.is_admin: - return - hub.register_platform_add_entities( - entity_class, descriptions, async_add_entities - ) - - @callback - def _async_should_add_entity( - self, description: UnifiEntityDescription, obj_id: str - ) -> bool: - """Check if entity should be added.""" - return bool( - (description.key, obj_id) not in self.known_objects - and description.allowed_fn(self, obj_id) - and description.supported_fn(self, obj_id) - ) - - @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(descriptions: Iterable[UnifiEntityDescription]) -> None: - """Load and subscribe to UniFi endpoints.""" - - @callback - def async_add_unifi_entities() -> None: - """Add UniFi entity.""" - async_add_entities( - unifi_platform_entity(obj_id, self, description) - for description in descriptions - for obj_id in description.api_handler_fn(self.api) - if self._async_should_add_entity(description, obj_id) - ) - - async_add_unifi_entities() - - @callback - def async_create_entity( - description: UnifiEntityDescription, event: ItemEvent, obj_id: str - ) -> None: - """Create new UniFi entity on event.""" - if self._async_should_add_entity(description, obj_id): - async_add_entities( - [unifi_platform_entity(obj_id, self, description)] - ) - - for description in descriptions: - description.api_handler_fn(self.api).subscribe( - partial(async_create_entity, description), ItemEvent.ADDED - ) - - self.config.entry.async_on_unload( - async_dispatcher_connect( - self.hass, - self.signal_options_update, - async_add_unifi_entities, - ) - ) - - async_load_entities(descriptions) - @property def signal_reachable(self) -> str: """Integration specific event to signal a change in connection status.""" diff --git a/homeassistant/components/unifi/image.py b/homeassistant/components/unifi/image.py index 3ea53d5b3f1..8360fb3b095 100644 --- a/homeassistant/components/unifi/image.py +++ b/homeassistant/components/unifi/image.py @@ -72,9 +72,7 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up image platform for UniFi Network integration.""" - UnifiHub.register_platform( - hass, - config_entry, + UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform( async_add_entities, UnifiImageEntity, ENTITY_DESCRIPTIONS, diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index ccdabc54d24..eda4599b02a 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -374,8 +374,8 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up sensors for UniFi Network integration.""" - UnifiHub.register_platform( - hass, config_entry, async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS + UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform( + async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS ) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 1e6e0beaa8f..92c25b90f06 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -321,9 +321,7 @@ async def async_setup_entry( ) -> None: """Set up switches for UniFi Network integration.""" async_update_unique_id(hass, config_entry) - UnifiHub.register_platform( - hass, - config_entry, + UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform( async_add_entities, UnifiSwitchEntity, ENTITY_DESCRIPTIONS, diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index a8a2dbe3b23..f112f47232c 100644 --- a/homeassistant/components/unifi/update.py +++ b/homeassistant/components/unifi/update.py @@ -76,9 +76,7 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up update entities for UniFi Network integration.""" - UnifiHub.register_platform( - hass, - config_entry, + UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform( async_add_entities, UnifiDeviceUpdateEntity, ENTITY_DESCRIPTIONS, diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index 0f8efc8b483..ee309ca2579 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/tests/components/unifi/test_config_flow.py @@ -76,14 +76,15 @@ DEVICES = [ ] WLANS = [ - {"_id": "1", "name": "SSID 1"}, + {"_id": "1", "name": "SSID 1", "enabled": True}, { "_id": "2", "name": "SSID 2", "name_combine_enabled": False, "name_combine_suffix": "_IOT", + "enabled": True, }, - {"_id": "3", "name": "SSID 4", "name_combine_enabled": False}, + {"_id": "3", "name": "SSID 4", "name_combine_enabled": False, "enabled": True}, ] DPI_GROUPS = [ @@ -542,12 +543,7 @@ async def test_simple_option_flow( ) -> None: """Test simple config flow options.""" config_entry = await setup_unifi_integration( - hass, - aioclient_mock, - clients_response=CLIENTS, - wlans_response=WLANS, - dpigroup_response=DPI_GROUPS, - dpiapp_response=[], + hass, aioclient_mock, clients_response=CLIENTS ) result = await hass.config_entries.options.async_init(