UniFi switch entity description (#81680)

* Consolidate switch entities to one class

* Move turn on/off into UnifiSwitchEntity

* Add event subscription
Remove storing entity for everything but legacy poe switch

* Only one entity class

* Improve generics naming

* Rename loader to description

* Improve control_fn naming

* Move wrongfully placed method that should only react to dpi apps being emptied

* Improve different methods

* Minor renaming and sorting

* Mark callbacks properly

homeassistant/components/unifi/diagnostics.py

@@ -95,7 +95,6 @@ async def async_get_config_entry_diagnostics(
         async_replace_dict_data(config_entry.as_dict(), macs_to_redact), REDACT_CONFIG
     )
     diag["site_role"] = controller.site_role
-    diag["entities"] = async_replace_dict_data(controller.entities, macs_to_redact)
     diag["clients"] = {
         macs_to_redact[k]: async_redact_data(
             async_replace_dict_data(v.raw, macs_to_redact), REDACT_CLIENTS

homeassistant/components/unifi/switch.py

@@ -7,24 +7,33 @@ Support for controlling deep packet inspection (DPI) restriction groups.
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
-from aiounifi.interfaces.api_handlers import ItemEvent
+import aiounifi
+from aiounifi.interfaces.api_handlers import CallbackType, ItemEvent, UnsubscribeType
 from aiounifi.interfaces.clients import Clients
 from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups
 from aiounifi.interfaces.outlets import Outlets
 from aiounifi.interfaces.ports import Ports
-from aiounifi.models.client import ClientBlockRequest
+from aiounifi.models.client import Client, ClientBlockRequest
 from aiounifi.models.device import (
     DeviceSetOutletRelayRequest,
     DeviceSetPoePortModeRequest,
 )
 from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest
+from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup
 from aiounifi.models.event import Event, EventKey
+from aiounifi.models.outlet import Outlet
+from aiounifi.models.port import Port
 
-from homeassistant.components.switch import DOMAIN, SwitchDeviceClass, SwitchEntity
+from homeassistant.components.switch import (
+    DOMAIN,
+    SwitchDeviceClass,
+    SwitchEntity,
+    SwitchEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
@@ -37,33 +46,219 @@ from homeassistant.helpers.entity import DeviceInfo, EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 
-from .const import (
+from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN, POE_SWITCH
-    ATTR_MANUFACTURER,
-    BLOCK_SWITCH,
-    DOMAIN as UNIFI_DOMAIN,
-    DPI_SWITCH,
-    OUTLET_SWITCH,
-    POE_SWITCH,
-)
 from .controller import UniFiController
 from .unifi_client import UniFiClient
 
 CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED)
 CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED)
 
-T = TypeVar("T")
+Data = TypeVar("Data")
+Handler = TypeVar("Handler")
+
+Subscription = Callable[[CallbackType, ItemEvent], UnsubscribeType]
+
+
+@callback
+def async_dpi_group_is_on_fn(
+    api: aiounifi.Controller, dpi_group: DPIRestrictionGroup
+) -> bool:
+    """Calculate if all apps are enabled."""
+    return all(
+        api.dpi_apps[app_id].enabled
+        for app_id in dpi_group.dpiapp_ids or []
+        if app_id in api.dpi_apps
+    )
+
+
+@callback
+def async_sub_device_available_fn(controller: UniFiController, obj_id: str) -> bool:
+    """Check if sub device object is disabled."""
+    device_id = obj_id.partition("_")[0]
+    device = controller.api.devices[device_id]
+    return controller.available and not device.disabled
+
+
+@callback
+def async_client_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
+    """Create device registry entry for client."""
+    client = api.clients[obj_id]
+    return DeviceInfo(
+        connections={(CONNECTION_NETWORK_MAC, obj_id)},
+        default_manufacturer=client.oui,
+        default_name=client.name or client.hostname,
+    )
+
+
+@callback
+def async_device_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
+    """Create device registry entry for device."""
+    if "_" in obj_id:  # Sub device
+        obj_id = obj_id.partition("_")[0]
+
+    device = api.devices[obj_id]
+    return DeviceInfo(
+        connections={(CONNECTION_NETWORK_MAC, device.mac)},
+        manufacturer=ATTR_MANUFACTURER,
+        model=device.model,
+        name=device.name or None,
+        sw_version=device.version,
+        hw_version=str(device.board_revision),
+    )
+
+
+@callback
+def async_dpi_group_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
+    """Create device registry entry for DPI group."""
+    return DeviceInfo(
+        entry_type=DeviceEntryType.SERVICE,
+        identifiers={(DOMAIN, f"unifi_controller_{obj_id}")},
+        manufacturer=ATTR_MANUFACTURER,
+        model="UniFi Network",
+        name="UniFi Network",
+    )
+
+
+async def async_block_client_control_fn(
+    api: aiounifi.Controller, obj_id: str, target: bool
+) -> None:
+    """Control network access of client."""
+    await api.request(ClientBlockRequest.create(obj_id, not target))
+
+
+async def async_dpi_group_control_fn(
+    api: aiounifi.Controller, obj_id: str, target: bool
+) -> None:
+    """Enable or disable DPI group."""
+    dpi_group = api.dpi_groups[obj_id]
+    await asyncio.gather(
+        *[
+            api.request(DPIRestrictionAppEnableRequest.create(app_id, target))
+            for app_id in dpi_group.dpiapp_ids or []
+        ]
+    )
+
+
+async def async_outlet_control_fn(
+    api: aiounifi.Controller, obj_id: str, target: bool
+) -> None:
+    """Control outlet relay."""
+    mac, _, index = obj_id.partition("_")
+    device = api.devices[mac]
+    await api.request(DeviceSetOutletRelayRequest.create(device, int(index), target))
+
+
+async def async_poe_port_control_fn(
+    api: aiounifi.Controller, obj_id: str, target: bool
+) -> None:
+    """Control poe state."""
+    mac, _, index = obj_id.partition("_")
+    device = api.devices[mac]
+    state = "auto" if target else "off"
+    await api.request(DeviceSetPoePortModeRequest.create(device, int(index), state))
 
 
 @dataclass
-class UnifiEntityLoader(Generic[T]):
+class UnifiEntityLoader(Generic[Handler, Data]):
     """Validate and load entities from different UniFi handlers."""
 
     allowed_fn: Callable[[UniFiController, str], bool]
-    entity_cls: type[UnifiBlockClientSwitch] | type[UnifiDPIRestrictionSwitch] | type[
+    api_handler_fn: Callable[[aiounifi.Controller], Handler]
-        UnifiOutletSwitch
+    available_fn: Callable[[UniFiController, str], bool]
-    ] | type[UnifiPoePortSwitch] | type[UnifiDPIRestrictionSwitch]
+    control_fn: Callable[[aiounifi.Controller, str, bool], Coroutine[Any, Any, None]]
-    handler_fn: Callable[[UniFiController], T]
+    device_info_fn: Callable[[aiounifi.Controller, str], DeviceInfo]
-    supported_fn: Callable[[T, str], bool | None]
+    event_is_on: tuple[EventKey, ...] | None
+    event_to_subscribe: tuple[EventKey, ...] | None
+    is_on_fn: Callable[[aiounifi.Controller, Data], bool]
+    name_fn: Callable[[Data], str | None]
+    object_fn: Callable[[aiounifi.Controller, str], Data]
+    supported_fn: Callable[[aiounifi.Controller, str], bool | None]
+    unique_id_fn: Callable[[str], str]
+
+
+@dataclass
+class UnifiEntityDescription(SwitchEntityDescription, UnifiEntityLoader[Handler, Data]):
+    """Class describing UniFi switch entity."""
+
+    custom_subscribe: Callable[[aiounifi.Controller], Subscription] | None = None
+
+
+ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = (
+    UnifiEntityDescription[Clients, Client](
+        key="Block client",
+        device_class=SwitchDeviceClass.SWITCH,
+        entity_category=EntityCategory.CONFIG,
+        has_entity_name=True,
+        icon="mdi:ethernet",
+        allowed_fn=lambda controller, obj_id: obj_id in controller.option_block_clients,
+        api_handler_fn=lambda api: api.clients,
+        available_fn=lambda controller, obj_id: controller.available,
+        control_fn=async_block_client_control_fn,
+        device_info_fn=async_client_device_info_fn,
+        event_is_on=CLIENT_UNBLOCKED,
+        event_to_subscribe=CLIENT_BLOCKED + CLIENT_UNBLOCKED,
+        is_on_fn=lambda api, client: not client.blocked,
+        name_fn=lambda client: None,
+        object_fn=lambda api, obj_id: api.clients[obj_id],
+        supported_fn=lambda api, obj_id: True,
+        unique_id_fn=lambda obj_id: f"block-{obj_id}",
+    ),
+    UnifiEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup](
+        key="DPI restriction",
+        entity_category=EntityCategory.CONFIG,
+        icon="mdi:network",
+        allowed_fn=lambda controller, obj_id: controller.option_dpi_restrictions,
+        api_handler_fn=lambda api: api.dpi_groups,
+        available_fn=lambda controller, obj_id: controller.available,
+        control_fn=async_dpi_group_control_fn,
+        custom_subscribe=lambda api: api.dpi_apps.subscribe,
+        device_info_fn=async_dpi_group_device_info_fn,
+        event_is_on=None,
+        event_to_subscribe=None,
+        is_on_fn=async_dpi_group_is_on_fn,
+        name_fn=lambda group: group.name,
+        object_fn=lambda api, obj_id: api.dpi_groups[obj_id],
+        supported_fn=lambda api, obj_id: bool(api.dpi_groups[obj_id].dpiapp_ids),
+        unique_id_fn=lambda obj_id: obj_id,
+    ),
+    UnifiEntityDescription[Outlets, Outlet](
+        key="Outlet control",
+        device_class=SwitchDeviceClass.OUTLET,
+        has_entity_name=True,
+        allowed_fn=lambda controller, obj_id: True,
+        api_handler_fn=lambda api: api.outlets,
+        available_fn=async_sub_device_available_fn,
+        control_fn=async_outlet_control_fn,
+        device_info_fn=async_device_device_info_fn,
+        event_is_on=None,
+        event_to_subscribe=None,
+        is_on_fn=lambda api, outlet: outlet.relay_state,
+        name_fn=lambda outlet: outlet.name,
+        object_fn=lambda api, obj_id: api.outlets[obj_id],
+        supported_fn=lambda api, obj_id: api.outlets[obj_id].has_relay,
+        unique_id_fn=lambda obj_id: f"{obj_id.split('_', 1)[0]}-outlet-{obj_id.split('_', 1)[1]}",
+    ),
+    UnifiEntityDescription[Ports, Port](
+        key="PoE port control",
+        device_class=SwitchDeviceClass.OUTLET,
+        entity_category=EntityCategory.CONFIG,
+        has_entity_name=True,
+        entity_registry_enabled_default=False,
+        icon="mdi:ethernet",
+        allowed_fn=lambda controller, obj_id: True,
+        api_handler_fn=lambda api: api.ports,
+        available_fn=async_sub_device_available_fn,
+        control_fn=async_poe_port_control_fn,
+        device_info_fn=async_device_device_info_fn,
+        event_is_on=None,
+        event_to_subscribe=None,
+        is_on_fn=lambda api, port: port.poe_mode != "off",
+        name_fn=lambda port: f"{port.name} PoE",
+        object_fn=lambda api, obj_id: api.ports[obj_id],
+        supported_fn=lambda api, obj_id: api.ports[obj_id].port_poe,
+        unique_id_fn=lambda obj_id: f"{obj_id.split('_', 1)[0]}-poe-{obj_id.split('_', 1)[1]}",
+    ),
+)
 
 
 async def async_setup_entry(
@@ -71,17 +266,9 @@ async def async_setup_entry(
     config_entry: ConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Set up switches for UniFi Network integration.
+    """Set up switches for UniFi Network integration."""
-
-    Switches are controlling network access and switch ports with POE.
-    """
     controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
-    controller.entities[DOMAIN] = {
+    controller.entities[DOMAIN] = {POE_SWITCH: set()}
-        BLOCK_SWITCH: set(),
-        POE_SWITCH: set(),
-        DPI_SWITCH: set(),
-        OUTLET_SWITCH: set(),
-    }
 
     if controller.site_role != "admin":
         return
@@ -125,20 +312,20 @@ async def async_setup_entry(
     known_poe_clients.clear()
 
     @callback
-    def async_load_entities(loader: UnifiEntityLoader) -> None:
+    def async_load_entities(description: UnifiEntityDescription) -> None:
         """Load and subscribe to UniFi devices."""
         entities: list[SwitchEntity] = []
-        api_handler = loader.handler_fn(controller)
+        api_handler = description.api_handler_fn(controller.api)
 
         @callback
         def async_create_entity(event: ItemEvent, obj_id: str) -> None:
             """Create UniFi entity."""
-            if not loader.allowed_fn(controller, obj_id) or not loader.supported_fn(
+            if not description.allowed_fn(
-                api_handler, obj_id
+                controller, obj_id
-            ):
+            ) or not description.supported_fn(controller.api, obj_id):
                 return
 
-            entity = loader.entity_cls(obj_id, controller)
+            entity = UnifiSwitchEntity(obj_id, controller, description)
             if event == ItemEvent.ADDED:
                 async_add_entities([entity])
                 return
@@ -150,8 +337,8 @@ async def async_setup_entry(
 
         api_handler.subscribe(async_create_entity, ItemEvent.ADDED)
 
-    for unifi_loader in UNIFI_LOADERS:
+    for description in ENTITY_DESCRIPTIONS:
-        async_load_entities(unifi_loader)
+        async_load_entities(description)
 
 
 @callback
@@ -301,51 +488,52 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity):
             await self.remove_item({self.client.mac})
 
 
-class UnifiBlockClientSwitch(SwitchEntity):
+class UnifiSwitchEntity(SwitchEntity):
-    """Representation of a blockable client."""
+    """Base representation of a UniFi switch."""
 
-    _attr_device_class = SwitchDeviceClass.SWITCH
+    entity_description: UnifiEntityDescription
-    _attr_entity_category = EntityCategory.CONFIG
-    _attr_has_entity_name = True
-    _attr_icon = "mdi:ethernet"
     _attr_should_poll = False
 
-    def __init__(self, obj_id: str, controller: UniFiController) -> None:
+    def __init__(
-        """Set up block switch."""
+        self,
-        controller.entities[DOMAIN][BLOCK_SWITCH].add(obj_id)
+        obj_id: str,
+        controller: UniFiController,
+        description: UnifiEntityDescription,
+    ) -> None:
+        """Set up UniFi switch entity."""
         self._obj_id = obj_id
         self.controller = controller
+        self.entity_description = description
 
         self._removed = False
 
-        client = controller.api.clients[obj_id]
+        self._attr_available = description.available_fn(controller, obj_id)
-        self._attr_available = controller.available
+        self._attr_device_info = description.device_info_fn(controller.api, obj_id)
-        self._attr_is_on = not client.blocked
+        self._attr_unique_id = description.unique_id_fn(obj_id)
-        self._attr_unique_id = f"{BLOCK_SWITCH}-{obj_id}"
+
-        self._attr_device_info = DeviceInfo(
+        obj = description.object_fn(self.controller.api, obj_id)
-            connections={(CONNECTION_NETWORK_MAC, obj_id)},
+        self._attr_is_on = description.is_on_fn(controller.api, obj)
-            default_manufacturer=client.oui,
+        self._attr_name = description.name_fn(obj)
-            default_name=client.name or client.hostname,
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn on switch."""
+        await self.entity_description.control_fn(
+            self.controller.api, self._obj_id, True
+        )
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn off switch."""
+        await self.entity_description.control_fn(
+            self.controller.api, self._obj_id, False
         )
 
     async def async_added_to_hass(self) -> None:
-        """Entity created."""
+        """Register callbacks."""
+        description = self.entity_description
+        handler = description.api_handler_fn(self.controller.api)
         self.async_on_remove(
-            self.controller.api.clients.subscribe(self.async_signalling_callback)
+            handler.subscribe(
-        )
+                self.async_signalling_callback,
-        self.async_on_remove(
-            self.controller.api.events.subscribe(
-                self.async_event_callback, CLIENT_BLOCKED + CLIENT_UNBLOCKED
-            )
-        )
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass, self.controller.signal_remove, self.remove_item
-            )
-        )
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass, self.controller.signal_options_update, self.options_updated
             )
         )
         self.async_on_remove(
@@ -355,31 +543,49 @@ class UnifiBlockClientSwitch(SwitchEntity):
                 self.async_signal_reachable_callback,
             )
         )
-
+        self.async_on_remove(
-    async def async_will_remove_from_hass(self) -> None:
+            async_dispatcher_connect(
-        """Disconnect object when removed."""
+                self.hass,
-        self.controller.entities[DOMAIN][BLOCK_SWITCH].remove(self._obj_id)
+                self.controller.signal_options_update,
+                self.options_updated,
+            )
+        )
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                self.controller.signal_remove,
+                self.remove_item,
+            )
+        )
+        if description.event_to_subscribe is not None:
+            self.async_on_remove(
+                self.controller.api.events.subscribe(
+                    self.async_event_callback,
+                    description.event_to_subscribe,
+                )
+            )
+        if description.custom_subscribe is not None:
+            self.async_on_remove(
+                description.custom_subscribe(self.controller.api)(
+                    self.async_signalling_callback, ItemEvent.CHANGED
+                ),
+            )
 
     @callback
     def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
-        """Update the clients state."""
+        """Update the switch state."""
-        if event == ItemEvent.DELETED:
+        if event == ItemEvent.DELETED and obj_id == self._obj_id:
             self.hass.async_create_task(self.remove_item({self._obj_id}))
             return
 
-        client = self.controller.api.clients[self._obj_id]
+        description = self.entity_description
-        self._attr_is_on = not client.blocked
+        if not description.supported_fn(self.controller.api, self._obj_id):
-        self._attr_available = self.controller.available
+            self.hass.async_create_task(self.remove_item({self._obj_id}))
-        self.async_write_ha_state()
-
-    @callback
-    def async_event_callback(self, event: Event) -> None:
-        """Event subscription callback."""
-        if event.mac != self._obj_id:
             return
-        if event.key in CLIENT_BLOCKED + CLIENT_UNBLOCKED:
+
-            self._attr_is_on = event.key in CLIENT_UNBLOCKED
+        obj = description.object_fn(self.controller.api, self._obj_id)
-        self._attr_available = self.controller.available
+        self._attr_is_on = description.is_on_fn(self.controller.api, obj)
+        self._attr_available = description.available_fn(self.controller, self._obj_id)
         self.async_write_ha_state()
 
     @callback
@@ -387,30 +593,28 @@ class UnifiBlockClientSwitch(SwitchEntity):
         """Call when controller connection state change."""
         self.async_signalling_callback(ItemEvent.ADDED, self._obj_id)
 
-    async def async_turn_on(self, **kwargs: Any) -> None:
+    @callback
-        """Turn on connectivity for client."""
+    def async_event_callback(self, event: Event) -> None:
-        await self.controller.api.request(
+        """Event subscription callback."""
-            ClientBlockRequest.create(self._obj_id, False)
+        if event.mac != self._obj_id:
-        )
+            return
 
-    async def async_turn_off(self, **kwargs: Any) -> None:
+        description = self.entity_description
-        """Turn off connectivity for client."""
+        assert isinstance(description.event_to_subscribe, tuple)
-        await self.controller.api.request(ClientBlockRequest.create(self._obj_id, True))
+        assert isinstance(description.event_is_on, tuple)
 
-    @property
+        if event.key in description.event_to_subscribe:
-    def icon(self) -> str:
+            self._attr_is_on = event.key in description.event_is_on
-        """Return the icon to use in the frontend."""
+        self._attr_available = description.available_fn(self.controller, self._obj_id)
-        if not self.is_on:
+        self.async_write_ha_state()
-            return "mdi:network-off"
-        return "mdi:network"
 
     async def options_updated(self) -> None:
         """Config entry options are updated, remove entity if option is disabled."""
-        if self._obj_id not in self.controller.option_block_clients:
+        if not self.entity_description.allowed_fn(self.controller, self._obj_id):
             await self.remove_item({self._obj_id})
 
     async def remove_item(self, keys: set) -> None:
-        """Remove entity if key is part of set."""
+        """Remove entity if object ID is part of set."""
         if self._obj_id not in keys or self._removed:
             return
         self._removed = True
@@ -418,313 +622,3 @@ class UnifiBlockClientSwitch(SwitchEntity):
             er.async_get(self.hass).async_remove(self.entity_id)
         else:
             await self.async_remove(force_remove=True)
-
-
-class UnifiDPIRestrictionSwitch(SwitchEntity):
-    """Representation of a DPI restriction group."""
-
-    _attr_entity_category = EntityCategory.CONFIG
-
-    def __init__(self, obj_id: str, controller: UniFiController) -> None:
-        """Set up dpi switch."""
-        controller.entities[DOMAIN][DPI_SWITCH].add(obj_id)
-        self._obj_id = obj_id
-        self.controller = controller
-
-        dpi_group = controller.api.dpi_groups[obj_id]
-        self._known_app_ids = dpi_group.dpiapp_ids
-
-        self._attr_available = controller.available
-        self._attr_is_on = self.calculate_enabled()
-        self._attr_name = dpi_group.name
-        self._attr_unique_id = dpi_group.id
-        self._attr_device_info = DeviceInfo(
-            entry_type=DeviceEntryType.SERVICE,
-            identifiers={(DOMAIN, f"unifi_controller_{obj_id}")},
-            manufacturer=ATTR_MANUFACTURER,
-            model="UniFi Network",
-            name="UniFi Network",
-        )
-
-    async def async_added_to_hass(self) -> None:
-        """Register callback to known apps."""
-        self.async_on_remove(
-            self.controller.api.dpi_groups.subscribe(self.async_signalling_callback)
-        )
-        self.async_on_remove(
-            self.controller.api.dpi_apps.subscribe(
-                self.async_signalling_callback, ItemEvent.CHANGED
-            ),
-        )
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass, self.controller.signal_remove, self.remove_item
-            )
-        )
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass, self.controller.signal_options_update, self.options_updated
-            )
-        )
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                self.controller.signal_reachable,
-                self.async_signal_reachable_callback,
-            )
-        )
-
-    async def async_will_remove_from_hass(self) -> None:
-        """Disconnect object when removed."""
-        self.controller.entities[DOMAIN][DPI_SWITCH].remove(self._obj_id)
-
-    @callback
-    def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
-        """Object has new event."""
-        if event == ItemEvent.DELETED:
-            self.hass.async_create_task(self.remove_item({self._obj_id}))
-            return
-
-        dpi_group = self.controller.api.dpi_groups[self._obj_id]
-        if not dpi_group.dpiapp_ids:
-            self.hass.async_create_task(self.remove_item({self._obj_id}))
-            return
-
-        self._attr_available = self.controller.available
-        self._attr_is_on = self.calculate_enabled()
-        self.async_write_ha_state()
-
-    @callback
-    def async_signal_reachable_callback(self) -> None:
-        """Call when controller connection state change."""
-        self.async_signalling_callback(ItemEvent.ADDED, self._obj_id)
-
-    @property
-    def icon(self):
-        """Return the icon to use in the frontend."""
-        if self.is_on:
-            return "mdi:network"
-        return "mdi:network-off"
-
-    def calculate_enabled(self) -> bool:
-        """Calculate if all apps are enabled."""
-        dpi_group = self.controller.api.dpi_groups[self._obj_id]
-        return all(
-            self.controller.api.dpi_apps[app_id].enabled
-            for app_id in dpi_group.dpiapp_ids
-            if app_id in self.controller.api.dpi_apps
-        )
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Restrict access of apps related to DPI group."""
-        dpi_group = self.controller.api.dpi_groups[self._obj_id]
-        return await asyncio.gather(
-            *[
-                self.controller.api.request(
-                    DPIRestrictionAppEnableRequest.create(app_id, True)
-                )
-                for app_id in dpi_group.dpiapp_ids
-            ]
-        )
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Remove restriction of apps related to DPI group."""
-        dpi_group = self.controller.api.dpi_groups[self._obj_id]
-        return await asyncio.gather(
-            *[
-                self.controller.api.request(
-                    DPIRestrictionAppEnableRequest.create(app_id, False)
-                )
-                for app_id in dpi_group.dpiapp_ids
-            ]
-        )
-
-    async def options_updated(self) -> None:
-        """Config entry options are updated, remove entity if option is disabled."""
-        if not self.controller.option_dpi_restrictions:
-            await self.remove_item({self._attr_unique_id})
-
-    async def remove_item(self, keys: set) -> None:
-        """Remove entity if key is part of set."""
-        if self._attr_unique_id not in keys:
-            return
-
-        if self.registry_entry:
-            er.async_get(self.hass).async_remove(self.entity_id)
-        else:
-            await self.async_remove(force_remove=True)
-
-
-class UnifiOutletSwitch(SwitchEntity):
-    """Representation of a outlet relay."""
-
-    _attr_device_class = SwitchDeviceClass.OUTLET
-    _attr_has_entity_name = True
-    _attr_should_poll = False
-
-    def __init__(self, obj_id: str, controller: UniFiController) -> None:
-        """Set up UniFi Network entity base."""
-        self._device_mac, index = obj_id.split("_", 1)
-        self._index = int(index)
-        self._obj_id = obj_id
-        self.controller = controller
-
-        outlet = self.controller.api.outlets[self._obj_id]
-        self._attr_name = outlet.name
-        self._attr_is_on = outlet.relay_state
-        self._attr_unique_id = f"{self._device_mac}-outlet-{index}"
-
-        device = self.controller.api.devices[self._device_mac]
-        self._attr_available = controller.available and not device.disabled
-        self._attr_device_info = DeviceInfo(
-            connections={(CONNECTION_NETWORK_MAC, device.mac)},
-            manufacturer=ATTR_MANUFACTURER,
-            model=device.model,
-            name=device.name or None,
-            sw_version=device.version,
-            hw_version=device.board_revision,
-        )
-
-    async def async_added_to_hass(self) -> None:
-        """Entity created."""
-        self.async_on_remove(
-            self.controller.api.outlets.subscribe(self.async_signalling_callback)
-        )
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                self.controller.signal_reachable,
-                self.async_signal_reachable_callback,
-            )
-        )
-
-    @callback
-    def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
-        """Object has new event."""
-        device = self.controller.api.devices[self._device_mac]
-        outlet = self.controller.api.outlets[self._obj_id]
-        self._attr_available = self.controller.available and not device.disabled
-        self._attr_is_on = outlet.relay_state
-        self.async_write_ha_state()
-
-    @callback
-    def async_signal_reachable_callback(self) -> None:
-        """Call when controller connection state change."""
-        self.async_signalling_callback(ItemEvent.ADDED, self._obj_id)
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Enable outlet relay."""
-        device = self.controller.api.devices[self._device_mac]
-        await self.controller.api.request(
-            DeviceSetOutletRelayRequest.create(device, self._index, True)
-        )
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Disable outlet relay."""
-        device = self.controller.api.devices[self._device_mac]
-        await self.controller.api.request(
-            DeviceSetOutletRelayRequest.create(device, self._index, False)
-        )
-
-
-class UnifiPoePortSwitch(SwitchEntity):
-    """Representation of a Power-over-Ethernet source port on an UniFi device."""
-
-    _attr_device_class = SwitchDeviceClass.OUTLET
-    _attr_entity_category = EntityCategory.CONFIG
-    _attr_entity_registry_enabled_default = False
-    _attr_has_entity_name = True
-    _attr_icon = "mdi:ethernet"
-    _attr_should_poll = False
-
-    def __init__(self, obj_id: str, controller: UniFiController) -> None:
-        """Set up UniFi Network entity base."""
-        self._device_mac, index = obj_id.split("_", 1)
-        self._index = int(index)
-        self._obj_id = obj_id
-        self.controller = controller
-
-        port = self.controller.api.ports[self._obj_id]
-        self._attr_name = f"{port.name} PoE"
-        self._attr_is_on = port.poe_mode != "off"
-        self._attr_unique_id = f"{self._device_mac}-poe-{index}"
-
-        device = self.controller.api.devices[self._device_mac]
-        self._attr_available = controller.available and not device.disabled
-        self._attr_device_info = DeviceInfo(
-            connections={(CONNECTION_NETWORK_MAC, device.mac)},
-            manufacturer=ATTR_MANUFACTURER,
-            model=device.model,
-            name=device.name or None,
-            sw_version=device.version,
-            hw_version=device.board_revision,
-        )
-
-    async def async_added_to_hass(self) -> None:
-        """Entity created."""
-        self.async_on_remove(
-            self.controller.api.ports.subscribe(self.async_signalling_callback)
-        )
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                self.controller.signal_reachable,
-                self.async_signal_reachable_callback,
-            )
-        )
-
-    @callback
-    def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
-        """Object has new event."""
-        device = self.controller.api.devices[self._device_mac]
-        port = self.controller.api.ports[self._obj_id]
-        self._attr_available = self.controller.available and not device.disabled
-        self._attr_is_on = port.poe_mode != "off"
-        self.async_write_ha_state()
-
-    @callback
-    def async_signal_reachable_callback(self) -> None:
-        """Call when controller connection state change."""
-        self.async_signalling_callback(ItemEvent.ADDED, self._obj_id)
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Enable POE for client."""
-        device = self.controller.api.devices[self._device_mac]
-        await self.controller.api.request(
-            DeviceSetPoePortModeRequest.create(device, self._index, "auto")
-        )
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Disable POE for client."""
-        device = self.controller.api.devices[self._device_mac]
-        await self.controller.api.request(
-            DeviceSetPoePortModeRequest.create(device, self._index, "off")
-        )
-
-
-UNIFI_LOADERS: tuple[UnifiEntityLoader, ...] = (
-    UnifiEntityLoader[Clients](
-        allowed_fn=lambda controller, obj_id: obj_id in controller.option_block_clients,
-        entity_cls=UnifiBlockClientSwitch,
-        handler_fn=lambda contrlr: contrlr.api.clients,
-        supported_fn=lambda handler, obj_id: True,
-    ),
-    UnifiEntityLoader[DPIRestrictionGroups](
-        allowed_fn=lambda controller, obj_id: controller.option_dpi_restrictions,
-        entity_cls=UnifiDPIRestrictionSwitch,
-        handler_fn=lambda controller: controller.api.dpi_groups,
-        supported_fn=lambda handler, obj_id: bool(handler[obj_id].dpiapp_ids),
-    ),
-    UnifiEntityLoader[Outlets](
-        allowed_fn=lambda controller, obj_id: True,
-        entity_cls=UnifiOutletSwitch,
-        handler_fn=lambda controller: controller.api.outlets,
-        supported_fn=lambda handler, obj_id: handler[obj_id].has_relay,
-    ),
-    UnifiEntityLoader[Ports](
-        allowed_fn=lambda controller, obj_id: True,
-        entity_cls=UnifiPoePortSwitch,
-        handler_fn=lambda controller: controller.api.ports,
-        supported_fn=lambda handler, obj_id: handler[obj_id].port_poe,
-    ),
-)

tests/components/unifi/test_diagnostics.py

@@ -6,16 +6,6 @@ from homeassistant.components.unifi.const import (
     CONF_ALLOW_UPTIME_SENSORS,
     CONF_BLOCK_CLIENT,
 )
-from homeassistant.components.unifi.device_tracker import CLIENT_TRACKER, DEVICE_TRACKER
-from homeassistant.components.unifi.sensor import RX_SENSOR, TX_SENSOR, UPTIME_SENSOR
-from homeassistant.components.unifi.switch import (
-    BLOCK_SWITCH,
-    DPI_SWITCH,
-    OUTLET_SWITCH,
-    POE_SWITCH,
-)
-from homeassistant.components.unifi.update import DEVICE_UPDATE
-from homeassistant.const import Platform
 
 from .test_controller import setup_unifi_integration
 
@@ -146,26 +136,6 @@ async def test_entry_diagnostics(hass, hass_client, aioclient_mock):
             "version": 1,
         },
         "site_role": "admin",
-        "entities": {
-            str(Platform.DEVICE_TRACKER): {
-                CLIENT_TRACKER: ["00:00:00:00:00:00"],
-                DEVICE_TRACKER: ["00:00:00:00:00:01"],
-            },
-            str(Platform.SENSOR): {
-                RX_SENSOR: ["00:00:00:00:00:00"],
-                TX_SENSOR: ["00:00:00:00:00:00"],
-                UPTIME_SENSOR: ["00:00:00:00:00:00"],
-            },
-            str(Platform.SWITCH): {
-                BLOCK_SWITCH: ["00:00:00:00:00:00"],
-                DPI_SWITCH: ["5f976f4ae3c58f018ec7dff6"],
-                POE_SWITCH: ["00:00:00:00:00:00"],
-                OUTLET_SWITCH: [],
-            },
-            str(Platform.UPDATE): {
-                DEVICE_UPDATE: ["00:00:00:00:00:01"],
-            },
-        },
         "clients": {
             "00:00:00:00:00:00": {
                 "blocked": False,