Refactor UniFi upgrade entities (#80752)

* Refactor UniFi upgrade entities

* Enable type check for UniFi update platform

.strict-typing

@@ -271,6 +271,7 @@ homeassistant.components.trafikverket_train.*
 homeassistant.components.trafikverket_weatherstation.*
 homeassistant.components.tts.*
 homeassistant.components.twentemilieu.*
+homeassistant.components.unifi.update
 homeassistant.components.unifiprotect.*
 homeassistant.components.upcloud.*
 homeassistant.components.update.*

homeassistant/components/unifi/update.py

@@ -2,8 +2,9 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
+from aiounifi.interfaces.api_handlers import ItemEvent
 from aiounifi.models.device import DeviceUpgradeRequest
 
 from homeassistant.components.update import (
@@ -13,7 +14,6 @@ from homeassistant.components.update import (
     UpdateEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import ATTR_NAME
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -21,7 +21,9 @@ from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
-from .unifi_entity_base import UniFiBase
+
+if TYPE_CHECKING:
+    from .controller import UniFiController
 
 LOGGER = logging.getLogger(__name__)
 
@@ -38,103 +40,84 @@ async def async_setup_entry(
     controller.entities[DOMAIN] = {DEVICE_UPDATE: set()}
 
     @callback
-    def items_added(
+    def async_add_update_entity(_: ItemEvent, obj_id: str) -> None:
-        clients: set = controller.api.clients, devices: set = controller.api.devices
-    ) -> None:
-        """Add device update entities."""
-        add_device_update_entities(controller, async_add_entities, devices)
-
-    for signal in (controller.signal_update, controller.signal_options_update):
-        config_entry.async_on_unload(
-            async_dispatcher_connect(hass, signal, items_added)
-        )
-
-    items_added()
-
-
-@callback
-def add_device_update_entities(controller, async_add_entities, devices):
         """Add new device update entities from the controller."""
-    entities = []
+        async_add_entities([UnifiDeviceUpdateEntity(obj_id, controller)])
 
-    for mac in devices:
+    controller.api.devices.subscribe(async_add_update_entity, ItemEvent.ADDED)
-        if mac in controller.entities[DOMAIN][UniFiDeviceUpdateEntity.TYPE]:
-            continue
 
-        device = controller.api.devices[mac]
+    for device_id in controller.api.devices:
-        entities.append(UniFiDeviceUpdateEntity(device, controller))
+        async_add_update_entity(ItemEvent.ADDED, device_id)
-
-    async_add_entities(entities)
 
 
-class UniFiDeviceUpdateEntity(UniFiBase, UpdateEntity):
+class UnifiDeviceUpdateEntity(UpdateEntity):
     """Update entity for a UniFi network infrastructure device."""
 
     DOMAIN = DOMAIN
     TYPE = DEVICE_UPDATE
     _attr_device_class = UpdateDeviceClass.FIRMWARE
+    _attr_has_entity_name = True
 
-    def __init__(self, device, controller):
+    def __init__(self, obj_id: str, controller: UniFiController) -> None:
         """Set up device update entity."""
-        super().__init__(device, controller)
+        controller.entities[DOMAIN][DEVICE_UPDATE].add(obj_id)
-
+        self.controller = controller
-        self.device = self._item
+        self._obj_id = obj_id
+        self._attr_unique_id = f"{self.TYPE}-{obj_id}"
 
         self._attr_supported_features = UpdateEntityFeature.PROGRESS
-
+        if controller.site_role == "admin":
-        if self.controller.site_role == "admin":
             self._attr_supported_features |= UpdateEntityFeature.INSTALL
 
-    @property
+        device = controller.api.devices[obj_id]
-    def name(self) -> str:
+        self._attr_available = controller.available and not device.disabled
-        """Return the name of the device."""
+        self._attr_in_progress = device.state == 4
-        return self.device.name or self.device.model
+        self._attr_installed_version = device.version
+        self._attr_latest_version = device.upgrade_to_firmware or device.version
 
-    @property
+        self._attr_device_info = DeviceInfo(
-    def unique_id(self) -> str:
+            connections={(CONNECTION_NETWORK_MAC, obj_id)},
-        """Return a unique identifier for this device."""
-        return f"{self.TYPE}-{self.device.mac}"
-
-    @property
-    def available(self) -> bool:
-        """Return if controller is available."""
-        return not self.device.disabled and self.controller.available
-
-    @property
-    def in_progress(self) -> bool:
-        """Update installation in progress."""
-        return self.device.state == 4
-
-    @property
-    def installed_version(self) -> str | None:
-        """Version currently in use."""
-        return self.device.version
-
-    @property
-    def latest_version(self) -> str | None:
-        """Latest version available for install."""
-        return self.device.upgrade_to_firmware or self.device.version
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return a device description for device registry."""
-        info = DeviceInfo(
-            connections={(CONNECTION_NETWORK_MAC, self.device.mac)},
             manufacturer=ATTR_MANUFACTURER,
-            model=self.device.model,
+            model=device.model,
-            sw_version=self.device.version,
+            name=device.name or None,
+            sw_version=device.version,
+            hw_version=device.board_revision,
         )
 
-        if self.device.name:
+    async def async_added_to_hass(self) -> None:
-            info[ATTR_NAME] = self.device.name
+        """Entity created."""
+        self.async_on_remove(
+            self.controller.api.devices.subscribe(self.async_signalling_callback)
+        )
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                self.controller.signal_reachable,
+                self.async_signal_reachable_callback,
+            )
+        )
 
-        return info
+    async def async_will_remove_from_hass(self) -> None:
+        """Disconnect object when removed."""
+        self.controller.entities[DOMAIN][DEVICE_UPDATE].remove(self._obj_id)
 
-    async def options_updated(self) -> None:
+    @callback
-        """No action needed."""
+    def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
+        """Object has new event."""
+        device = self.controller.api.devices[self._obj_id]
+        self._attr_available = self.controller.available and not device.disabled
+        self._attr_in_progress = device.state == 4
+        self._attr_installed_version = device.version
+        self._attr_latest_version = device.upgrade_to_firmware or device.version
+        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_install(
         self, version: str | None, backup: bool, **kwargs: Any
     ) -> None:
         """Install an update."""
-        await self.controller.api.request(DeviceUpgradeRequest.create(self.device.mac))
+        await self.controller.api.request(DeviceUpgradeRequest.create(self._obj_id))

mypy.ini

@@ -2463,6 +2463,16 @@ disallow_untyped_defs = true
 warn_return_any = true
 warn_unreachable = true
 
+[mypy-homeassistant.components.unifi.update]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+warn_return_any = true
+warn_unreachable = true
+
 [mypy-homeassistant.components.unifiprotect.*]
 check_untyped_defs = true
 disallow_incomplete_defs = true