Align implementation of the UniFi update platform with the UniFi switch platform (#81821)

* Use the same entity class as switches, but separately

Once all platforms are migrated I will consolidate them into one entity class

* Fix review comment

* Fix review comments

* Fix review comments

homeassistant/components/unifi/switch.py

@@ -9,14 +9,20 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, TypeVar, Union
 
 import aiounifi
-from aiounifi.interfaces.api_handlers import CallbackType, ItemEvent, UnsubscribeType
+from aiounifi.interfaces.api_handlers import (
+    APIHandler,
+    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.api import APIItem
 from aiounifi.models.client import Client, ClientBlockRequest
 from aiounifi.models.device import (
     DeviceSetOutletRelayRequest,
@@ -51,8 +57,8 @@ from .controller import UniFiController
 CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED)
 CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED)
 
-Data = TypeVar("Data")
+_DataT = TypeVar("_DataT", bound=Union[APIItem, Outlet, Port])
-Handler = TypeVar("Handler")
+_HandlerT = TypeVar("_HandlerT", bound=Union[APIHandler, Outlets, Ports])
 
 Subscription = Callable[[CallbackType, ItemEvent], UnsubscribeType]
 
@@ -157,25 +163,27 @@ async def async_poe_port_control_fn(
 
 
 @dataclass
-class UnifiEntityLoader(Generic[Handler, Data]):
+class UnifiEntityLoader(Generic[_HandlerT, _DataT]):
     """Validate and load entities from different UniFi handlers."""
 
     allowed_fn: Callable[[UniFiController, str], bool]
-    api_handler_fn: Callable[[aiounifi.Controller], Handler]
+    api_handler_fn: Callable[[aiounifi.Controller], _HandlerT]
     available_fn: Callable[[UniFiController, str], bool]
     control_fn: Callable[[aiounifi.Controller, str, bool], Coroutine[Any, Any, None]]
     device_info_fn: Callable[[aiounifi.Controller, str], DeviceInfo]
     event_is_on: tuple[EventKey, ...] | None
     event_to_subscribe: tuple[EventKey, ...] | None
-    is_on_fn: Callable[[aiounifi.Controller, Data], bool]
+    is_on_fn: Callable[[aiounifi.Controller, _DataT], bool]
-    name_fn: Callable[[Data], str | None]
+    name_fn: Callable[[_DataT], str | None]
-    object_fn: Callable[[aiounifi.Controller, str], Data]
+    object_fn: Callable[[aiounifi.Controller, str], _DataT]
     supported_fn: Callable[[aiounifi.Controller, str], bool | None]
     unique_id_fn: Callable[[str], str]
 
 
 @dataclass
-class UnifiEntityDescription(SwitchEntityDescription, UnifiEntityLoader[Handler, Data]):
+class UnifiEntityDescription(
+    SwitchEntityDescription, UnifiEntityLoader[_HandlerT, _DataT]
+):
     """Class describing UniFi switch entity."""
 
     custom_subscribe: Callable[[aiounifi.Controller], Subscription] | None = None
@@ -307,17 +315,17 @@ async def async_setup_entry(
         async_load_entities(description)
 
 
-class UnifiSwitchEntity(SwitchEntity):
+class UnifiSwitchEntity(SwitchEntity, Generic[_HandlerT, _DataT]):
     """Base representation of a UniFi switch."""
 
-    entity_description: UnifiEntityDescription
+    entity_description: UnifiEntityDescription[_HandlerT, _DataT]
     _attr_should_poll = False
 
     def __init__(
         self,
         obj_id: str,
         controller: UniFiController,
-        description: UnifiEntityDescription,
+        description: UnifiEntityDescription[_HandlerT, _DataT],
     ) -> None:
         """Set up UniFi switch entity."""
         self._obj_id = obj_id

homeassistant/components/unifi/update.py

@@ -1,20 +1,25 @@
 """Update entities for Ubiquiti network devices."""
 from __future__ import annotations
 
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
 import logging
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
-from aiounifi.interfaces.api_handlers import ItemEvent
+import aiounifi
-from aiounifi.models.device import DeviceUpgradeRequest
+from aiounifi.interfaces.api_handlers import CallbackType, ItemEvent, UnsubscribeType
+from aiounifi.interfaces.devices import Devices
+from aiounifi.models.device import Device, DeviceUpgradeRequest
 
 from homeassistant.components.update import (
-    DOMAIN,
     UpdateDeviceClass,
     UpdateEntity,
+    UpdateEntityDescription,
     UpdateEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import DeviceInfo
@@ -23,11 +28,90 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
 
 if TYPE_CHECKING:
+    from aiounifi.models.event import EventKey
+
     from .controller import UniFiController
 
+_DataT = TypeVar("_DataT", bound=Device)
+_HandlerT = TypeVar("_HandlerT", bound=Devices)
+
+Subscription = Callable[[CallbackType, ItemEvent], UnsubscribeType]
+
 LOGGER = logging.getLogger(__name__)
 
-DEVICE_UPDATE = "device_update"
+
+@callback
+def async_device_available_fn(controller: UniFiController, obj_id: str) -> bool:
+    """Check if device is available."""
+    device = controller.api.devices[obj_id]
+    return controller.available and not device.disabled
+
+
+async def async_device_control_fn(api: aiounifi.Controller, obj_id: str) -> None:
+    """Control upgrade of device."""
+    await api.request(DeviceUpgradeRequest.create(obj_id))
+
+
+@callback
+def async_device_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
+    """Create device registry entry for device."""
+    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),
+    )
+
+
+@dataclass
+class UnifiEntityLoader(Generic[_HandlerT, _DataT]):
+    """Validate and load entities from different UniFi handlers."""
+
+    allowed_fn: Callable[[UniFiController, str], bool]
+    api_handler_fn: Callable[[aiounifi.Controller], _HandlerT]
+    available_fn: Callable[[UniFiController, str], bool]
+    control_fn: Callable[[aiounifi.Controller, str], Coroutine[Any, Any, None]]
+    device_info_fn: Callable[[aiounifi.Controller, str], DeviceInfo]
+    event_is_on: tuple[EventKey, ...] | None
+    event_to_subscribe: tuple[EventKey, ...] | None
+    name_fn: Callable[[_DataT], str | None]
+    object_fn: Callable[[aiounifi.Controller, str], _DataT]
+    state_fn: Callable[[aiounifi.Controller, _DataT], bool]
+    supported_fn: Callable[[aiounifi.Controller, str], bool | None]
+    unique_id_fn: Callable[[str], str]
+
+
+@dataclass
+class UnifiEntityDescription(
+    UpdateEntityDescription, UnifiEntityLoader[_HandlerT, _DataT]
+):
+    """Class describing UniFi update entity."""
+
+    custom_subscribe: Callable[[aiounifi.Controller], Subscription] | None = None
+
+
+ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = (
+    UnifiEntityDescription[Devices, Device](
+        key="Upgrade device",
+        device_class=UpdateDeviceClass.FIRMWARE,
+        has_entity_name=True,
+        allowed_fn=lambda controller, obj_id: True,
+        api_handler_fn=lambda api: api.devices,
+        available_fn=async_device_available_fn,
+        control_fn=async_device_control_fn,
+        device_info_fn=async_device_device_info_fn,
+        event_is_on=None,
+        event_to_subscribe=None,
+        name_fn=lambda device: None,
+        object_fn=lambda api, obj_id: api.devices[obj_id],
+        state_fn=lambda api, device: device.state == 4,
+        supported_fn=lambda api, obj_id: True,
+        unique_id_fn=lambda obj_id: f"device_update-{obj_id}",
+    ),
+)
 
 
 async def async_setup_entry(
@@ -36,58 +120,85 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up update entities for UniFi Network integration."""
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+    controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
-    controller.entities[DOMAIN] = {DEVICE_UPDATE: set()}
 
     @callback
-    def async_add_update_entity(_: ItemEvent, obj_id: str) -> None:
+    def async_load_entities(description: UnifiEntityDescription) -> None:
-        """Add new device update entities from the controller."""
+        """Load and subscribe to UniFi devices."""
-        async_add_entities([UnifiDeviceUpdateEntity(obj_id, controller)])
+        entities: list[UpdateEntity] = []
+        api_handler = description.api_handler_fn(controller.api)
 
-    controller.api.devices.subscribe(async_add_update_entity, ItemEvent.ADDED)
+        @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
 
-    for device_id in controller.api.devices:
+            entity = UnifiDeviceUpdateEntity(obj_id, controller, description)
-        async_add_update_entity(ItemEvent.ADDED, device_id)
+            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)
 
 
-class UnifiDeviceUpdateEntity(UpdateEntity):
+class UnifiDeviceUpdateEntity(UpdateEntity, Generic[_HandlerT, _DataT]):
-    """Update entity for a UniFi network infrastructure device."""
+    """Representation of a UniFi device update entity."""
 
-    DOMAIN = DOMAIN
+    entity_description: UnifiEntityDescription[_HandlerT, _DataT]
-    TYPE = DEVICE_UPDATE
+    _attr_should_poll = False
-    _attr_device_class = UpdateDeviceClass.FIRMWARE
-    _attr_has_entity_name = True
 
-    def __init__(self, obj_id: str, controller: UniFiController) -> None:
+    def __init__(
-        """Set up device update entity."""
+        self,
-        controller.entities[DOMAIN][DEVICE_UPDATE].add(obj_id)
+        obj_id: str,
-        self.controller = controller
+        controller: UniFiController,
+        description: UnifiEntityDescription[_HandlerT, _DataT],
+    ) -> None:
+        """Set up UniFi update entity."""
         self._obj_id = obj_id
-        self._attr_unique_id = f"{self.TYPE}-{obj_id}"
+        self.controller = controller
+        self.entity_description = description
+
+        self._removed = False
 
         self._attr_supported_features = UpdateEntityFeature.PROGRESS
         if controller.site_role == "admin":
             self._attr_supported_features |= UpdateEntityFeature.INSTALL
 
-        device = controller.api.devices[obj_id]
+        self._attr_available = description.available_fn(controller, obj_id)
-        self._attr_available = controller.available and not device.disabled
+        self._attr_device_info = description.device_info_fn(controller.api, obj_id)
-        self._attr_in_progress = device.state == 4
+        self._attr_unique_id = description.unique_id_fn(obj_id)
-        self._attr_installed_version = device.version
-        self._attr_latest_version = device.upgrade_to_firmware or device.version
 
-        self._attr_device_info = DeviceInfo(
+        obj = description.object_fn(self.controller.api, obj_id)
-            connections={(CONNECTION_NETWORK_MAC, obj_id)},
+        self._attr_in_progress = description.state_fn(controller.api, obj)
-            manufacturer=ATTR_MANUFACTURER,
+        self._attr_name = description.name_fn(obj)
-            model=device.model,
+        self._attr_installed_version = obj.version
-            name=device.name or None,
+        self._attr_latest_version = obj.upgrade_to_firmware or obj.version
-            sw_version=device.version,
+
-            hw_version=device.board_revision,
+    async def async_install(
-        )
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Install an update."""
+        await self.entity_description.control_fn(self.controller.api, self._obj_id)
 
     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.devices.subscribe(self.async_signalling_callback)
+            handler.subscribe(
+                self.async_signalling_callback,
+            )
         )
         self.async_on_remove(
             async_dispatcher_connect(
@@ -96,19 +207,27 @@ class UnifiDeviceUpdateEntity(UpdateEntity):
                 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][DEVICE_UPDATE].remove(self._obj_id)
+                self.controller.signal_remove,
+                self.remove_item,
+            )
+        )
 
     @callback
     def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
-        """Object has new event."""
+        """Update the switch state."""
-        device = self.controller.api.devices[self._obj_id]
+        if event == ItemEvent.DELETED and obj_id == self._obj_id:
-        self._attr_available = self.controller.available and not device.disabled
+            self.hass.async_create_task(self.remove_item({self._obj_id}))
-        self._attr_in_progress = device.state == 4
+            return
-        self._attr_installed_version = device.version
+
-        self._attr_latest_version = device.upgrade_to_firmware or device.version
+        description = self.entity_description
+        obj = description.object_fn(self.controller.api, self._obj_id)
+        self._attr_available = description.available_fn(self.controller, self._obj_id)
+        self._attr_in_progress = description.state_fn(self.controller.api, obj)
+        self._attr_installed_version = obj.version
+        self._attr_latest_version = obj.upgrade_to_firmware or obj.version
         self.async_write_ha_state()
 
     @callback
@@ -116,8 +235,12 @@ class UnifiDeviceUpdateEntity(UpdateEntity):
         """Call when controller connection state change."""
         self.async_signalling_callback(ItemEvent.ADDED, self._obj_id)
 
-    async def async_install(
+    async def remove_item(self, keys: set) -> None:
-        self, version: str | None, backup: bool, **kwargs: Any
+        """Remove entity if object ID is part of set."""
-    ) -> None:
+        if self._obj_id not in keys or self._removed:
-        """Install an update."""
+            return
-        await self.controller.api.request(DeviceUpgradeRequest.create(self._obj_id))
+        self._removed = True
+        if self.registry_entry:
+            er.async_get(self.hass).async_remove(self.entity_id)
+        else:
+            await self.async_remove(force_remove=True)