From 188ce9bf49c77b5056f3424a32210a72765ab4ac Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Sun, 11 Dec 2022 16:41:58 +0100 Subject: [PATCH] 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 | 34 ++-- homeassistant/components/unifi/update.py | 231 +++++++++++++++++------ 2 files changed, 198 insertions(+), 67 deletions(-) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 3e98f41188a..44007e4c1a8 100644 --- a/homeassistant/components/unifi/switch.py +++ b/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") -Handler = TypeVar("Handler") +_DataT = TypeVar("_DataT", bound=Union[APIItem, Outlet, Port]) +_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] - name_fn: Callable[[Data], str | None] - object_fn: Callable[[aiounifi.Controller, str], Data] + is_on_fn: Callable[[aiounifi.Controller, _DataT], bool] + name_fn: Callable[[_DataT], str | None] + 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 diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index 5f26cba57cd..4cd5282f047 100644 --- a/homeassistant/components/unifi/update.py +++ b/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 -from aiounifi.models.device import DeviceUpgradeRequest +import aiounifi +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.entities[DOMAIN] = {DEVICE_UPDATE: set()} + controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] @callback - def async_add_update_entity(_: ItemEvent, obj_id: str) -> None: - """Add new device update entities from the controller.""" - async_add_entities([UnifiDeviceUpdateEntity(obj_id, controller)]) + def async_load_entities(description: UnifiEntityDescription) -> None: + """Load and subscribe to UniFi devices.""" + 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: - async_add_update_entity(ItemEvent.ADDED, device_id) + entity = UnifiDeviceUpdateEntity(obj_id, controller, description) + 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): - """Update entity for a UniFi network infrastructure device.""" +class UnifiDeviceUpdateEntity(UpdateEntity, Generic[_HandlerT, _DataT]): + """Representation of a UniFi device update entity.""" - DOMAIN = DOMAIN - TYPE = DEVICE_UPDATE - _attr_device_class = UpdateDeviceClass.FIRMWARE - _attr_has_entity_name = True + entity_description: UnifiEntityDescription[_HandlerT, _DataT] + _attr_should_poll = False - def __init__(self, obj_id: str, controller: UniFiController) -> None: - """Set up device update entity.""" - controller.entities[DOMAIN][DEVICE_UPDATE].add(obj_id) - self.controller = controller + def __init__( + self, + obj_id: str, + 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 = 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._attr_available = description.available_fn(controller, obj_id) + self._attr_device_info = description.device_info_fn(controller.api, obj_id) + self._attr_unique_id = description.unique_id_fn(obj_id) - self._attr_device_info = DeviceInfo( - connections={(CONNECTION_NETWORK_MAC, obj_id)}, - manufacturer=ATTR_MANUFACTURER, - model=device.model, - name=device.name or None, - sw_version=device.version, - hw_version=device.board_revision, - ) + obj = description.object_fn(self.controller.api, obj_id) + self._attr_in_progress = description.state_fn(controller.api, obj) + self._attr_name = description.name_fn(obj) + self._attr_installed_version = obj.version + self._attr_latest_version = obj.upgrade_to_firmware or obj.version + + 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, ) ) - - async def async_will_remove_from_hass(self) -> None: - """Disconnect object when removed.""" - self.controller.entities[DOMAIN][DEVICE_UPDATE].remove(self._obj_id) + self.async_on_remove( + async_dispatcher_connect( + self.hass, + self.controller.signal_remove, + self.remove_item, + ) + ) @callback 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 + """Update the switch state.""" + if event == ItemEvent.DELETED and obj_id == self._obj_id: + self.hass.async_create_task(self.remove_item({self._obj_id})) + return + + 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( - self, version: str | None, backup: bool, **kwargs: Any - ) -> None: - """Install an update.""" - await self.controller.api.request(DeviceUpgradeRequest.create(self._obj_id)) + async def remove_item(self, keys: set) -> None: + """Remove entity if object ID is part of set.""" + if self._obj_id not in keys or self._removed: + return + 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)