Break out UniFi platform registration to its own class (#112514)

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)

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(
+    UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
-        hass,
-        config_entry,
         async_add_entities,
         UnifiButtonEntity,
         ENTITY_DESCRIPTIONS,

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(
+    UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
-        hass, config_entry, async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS
+        async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS
     )
 
 

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)
 

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)

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 (
+from homeassistant.helpers.dispatcher import async_dispatcher_send
-    async_dispatcher_connect,
-    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_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."""

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(
+    UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
-        hass,
-        config_entry,
         async_add_entities,
         UnifiImageEntity,
         ENTITY_DESCRIPTIONS,

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(
+    UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
-        hass, config_entry, async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS
+        async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS
     )
 
 

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(
+    UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
-        hass,
-        config_entry,
         async_add_entities,
         UnifiSwitchEntity,
         ENTITY_DESCRIPTIONS,

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(
+    UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
-        hass,
-        config_entry,
         async_add_entities,
         UnifiDeviceUpdateEntity,
         ENTITY_DESCRIPTIONS,

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,
+        hass, aioclient_mock, clients_response=CLIENTS
-        aioclient_mock,
-        clients_response=CLIENTS,
-        wlans_response=WLANS,
-        dpigroup_response=DPI_GROUPS,
-        dpiapp_response=[],
     )
 
     result = await hass.config_entries.options.async_init(