From 30df9e7706270daa84560746a3cceb74ee095985 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Thu, 22 Dec 2022 13:17:17 +0100 Subject: [PATCH] Split out part of update sensor to a common UniFi entity class (#84262) * Split out part of update sensor to a common entity class * Mark methods abstract and/or raise NotImplementedError * Resolve review comments --- homeassistant/components/unifi/entity.py | 157 +++++++++++++++++++++++ homeassistant/components/unifi/update.py | 130 ++++--------------- 2 files changed, 183 insertions(+), 104 deletions(-) create mode 100644 homeassistant/components/unifi/entity.py diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py new file mode 100644 index 00000000000..da9ca3828d7 --- /dev/null +++ b/homeassistant/components/unifi/entity.py @@ -0,0 +1,157 @@ +"""UniFi entity representation.""" +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Callable +from dataclasses import dataclass +from typing import Generic, TypeVar + +import aiounifi +from aiounifi.interfaces.api_handlers import CallbackType, ItemEvent, UnsubscribeType +from aiounifi.interfaces.devices import Devices +from aiounifi.models.device import Device +from aiounifi.models.event import EventKey + +from homeassistant.core import callback +from homeassistant.helpers import entity_registry as er +from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity import DeviceInfo, Entity, EntityDescription + +from .controller import UniFiController + +DataT = TypeVar("DataT", bound=Device) +HandlerT = TypeVar("HandlerT", bound=Devices) +SubscriptionT = Callable[[CallbackType, ItemEvent], UnsubscribeType] + + +@dataclass +class UnifiDescription(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] + device_info_fn: Callable[[aiounifi.Controller, str], DeviceInfo | None] + 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] + supported_fn: Callable[[UniFiController, str], bool | None] + unique_id_fn: Callable[[UniFiController, str], str] + + +@dataclass +class UnifiEntityDescription(EntityDescription, UnifiDescription[HandlerT, DataT]): + """UniFi Entity Description.""" + + +class UnifiEntity(Entity, Generic[HandlerT, DataT]): + """Representation of a UniFi entity.""" + + entity_description: UnifiEntityDescription[HandlerT, DataT] + _attr_should_poll = False + + _attr_unique_id: str + + def __init__( + self, + obj_id: str, + controller: UniFiController, + description: UnifiEntityDescription[HandlerT, DataT], + ) -> None: + """Set up UniFi switch entity.""" + self._obj_id = obj_id + self.controller = controller + self.entity_description = description + + self._removed = False + + 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(controller, obj_id) + + obj = description.object_fn(self.controller.api, obj_id) + self._attr_name = description.name_fn(obj) + self.async_initiate_state() + + async def async_added_to_hass(self) -> None: + """Register callbacks.""" + description = self.entity_description + handler = description.api_handler_fn(self.controller.api) + + # New data from handler + self.async_on_remove( + handler.subscribe( + self.async_signalling_callback, + ) + ) + + # State change from controller or websocket + self.async_on_remove( + async_dispatcher_connect( + self.hass, + self.controller.signal_reachable, + self.async_signal_reachable_callback, + ) + ) + + # Config entry options updated + self.async_on_remove( + async_dispatcher_connect( + self.hass, + self.controller.signal_options_update, + self.async_signal_options_updated, + ) + ) + + @callback + def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: + """Update the entity 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 + if not description.supported_fn(self.controller, self._obj_id): + self.hass.async_create_task(self.remove_item({self._obj_id})) + return + + self._attr_available = description.available_fn(self.controller, self._obj_id) + self.async_update_state(event, obj_id) + 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_signal_options_updated(self) -> None: + """Config entry options are updated, remove entity if option is disabled.""" + 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 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) + + @callback + @abstractmethod + def async_initiate_state(self) -> None: + """Initiate entity state. + + Perform additional actions setting up platform entity child class state. + """ + + @callback + @abstractmethod + def async_update_state(self, event: ItemEvent, obj_id: str) -> None: + """Update entity state. + + Perform additional actions updating platform entity child class state. + """ diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index 4cd5282f047..6cff6b7932d 100644 --- a/homeassistant/components/unifi/update.py +++ b/homeassistant/components/unifi/update.py @@ -4,10 +4,10 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from dataclasses import dataclass import logging -from typing import TYPE_CHECKING, Any, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Generic import aiounifi -from aiounifi.interfaces.api_handlers import CallbackType, ItemEvent, UnsubscribeType +from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.devices import Devices from aiounifi.models.device import Device, DeviceUpgradeRequest @@ -19,24 +19,16 @@ from homeassistant.components.update import ( ) 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 from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN +from .entity import DataT, HandlerT, UnifiEntity, UnifiEntityDescription 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__) @@ -67,34 +59,24 @@ def async_device_device_info_fn(api: aiounifi.Controller, obj_id: str) -> Device @dataclass -class UnifiEntityLoader(Generic[_HandlerT, _DataT]): +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] + state_fn: Callable[[aiounifi.Controller, DataT], bool] @dataclass -class UnifiEntityDescription( - UpdateEntityDescription, UnifiEntityLoader[_HandlerT, _DataT] +class UnifiUpdateEntityDescription( + UpdateEntityDescription, + UnifiEntityDescription[HandlerT, DataT], + UnifiEntityLoader[HandlerT, DataT], ): """Class describing UniFi update entity.""" - custom_subscribe: Callable[[aiounifi.Controller], Subscription] | None = None - -ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( - UnifiEntityDescription[Devices, Device]( +ENTITY_DESCRIPTIONS: tuple[UnifiUpdateEntityDescription, ...] = ( + UnifiUpdateEntityDescription[Devices, Device]( key="Upgrade device", device_class=UpdateDeviceClass.FIRMWARE, has_entity_name=True, @@ -108,8 +90,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( 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}", + supported_fn=lambda controller, obj_id: True, + unique_id_fn=lambda controller, obj_id: f"device_update-{obj_id}", ), ) @@ -123,7 +105,7 @@ async def async_setup_entry( controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] @callback - def async_load_entities(description: UnifiEntityDescription) -> None: + def async_load_entities(description: UnifiUpdateEntityDescription) -> None: """Load and subscribe to UniFi devices.""" entities: list[UpdateEntity] = [] api_handler = description.api_handler_fn(controller.api) @@ -152,38 +134,19 @@ async def async_setup_entry( async_load_entities(description) -class UnifiDeviceUpdateEntity(UpdateEntity, Generic[_HandlerT, _DataT]): +class UnifiDeviceUpdateEntity(UnifiEntity[HandlerT, DataT], UpdateEntity): """Representation of a UniFi device update entity.""" - entity_description: UnifiEntityDescription[_HandlerT, _DataT] - _attr_should_poll = False - - def __init__( - self, - obj_id: str, - controller: UniFiController, - description: UnifiEntityDescription[_HandlerT, _DataT], - ) -> None: - """Set up UniFi update entity.""" - self._obj_id = obj_id - self.controller = controller - self.entity_description = description - - self._removed = False + entity_description: UnifiUpdateEntityDescription[HandlerT, DataT] + @callback + def async_initiate_state(self) -> None: + """Initiate entity state.""" self._attr_supported_features = UpdateEntityFeature.PROGRESS - if controller.site_role == "admin": + if self.controller.site_role == "admin": self._attr_supported_features |= UpdateEntityFeature.INSTALL - 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) - - 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 + self.async_update_state(ItemEvent.ADDED, self._obj_id) async def async_install( self, version: str | None, backup: bool, **kwargs: Any @@ -191,56 +154,15 @@ class UnifiDeviceUpdateEntity(UpdateEntity, Generic[_HandlerT, _DataT]): """Install an update.""" await self.entity_description.control_fn(self.controller.api, self._obj_id) - async def async_added_to_hass(self) -> None: - """Register callbacks.""" - description = self.entity_description - handler = description.api_handler_fn(self.controller.api) - self.async_on_remove( - handler.subscribe( - self.async_signalling_callback, - ) - ) - self.async_on_remove( - async_dispatcher_connect( - self.hass, - self.controller.signal_reachable, - self.async_signal_reachable_callback, - ) - ) - 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: - """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 + def async_update_state(self, event: ItemEvent, obj_id: str) -> None: + """Update entity state. + Update in_progress, installed_version and latest_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 - def async_signal_reachable_callback(self) -> None: - """Call when controller connection state change.""" - self.async_signalling_callback(ItemEvent.ADDED, 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)