From 6718b401815eda6fff67ece5cffa5b09ce33942f Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 3 Jan 2023 22:57:44 +0100 Subject: [PATCH] Make switch platform use common UniFi entity class (#84458) * Make switch platform use common UniFi entity class * Consolidate common functions between update and switch platforms * Use controller.register_platform_add_entities * Rename UnfiEntityLoader to UnifiUpdateEntityDescriptionMixin --- homeassistant/components/unifi/entity.py | 69 +++++- homeassistant/components/unifi/switch.py | 269 ++++++----------------- homeassistant/components/unifi/update.py | 47 ++-- 3 files changed, 147 insertions(+), 238 deletions(-) diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index b7aed362133..ff6d368bf97 100644 --- a/homeassistant/components/unifi/entity.py +++ b/homeassistant/components/unifi/entity.py @@ -4,27 +4,65 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING, Generic, TypeVar, Union 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 aiounifi.interfaces.api_handlers import ( + APIHandler, + CallbackType, + ItemEvent, + UnsubscribeType, +) +from aiounifi.interfaces.outlets import Outlets +from aiounifi.interfaces.ports import Ports +from aiounifi.models.api import APIItem +from aiounifi.models.event import Event, EventKey +from aiounifi.models.outlet import Outlet +from aiounifi.models.port import Port from homeassistant.core import 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, Entity, EntityDescription +from .const import ATTR_MANUFACTURER + if TYPE_CHECKING: from .controller import UniFiController -DataT = TypeVar("DataT", bound=Device) -HandlerT = TypeVar("HandlerT", bound=Devices) +DataT = TypeVar("DataT", bound=Union[APIItem, Outlet, Port]) +HandlerT = TypeVar("HandlerT", bound=Union[APIHandler, Outlets, Ports]) SubscriptionT = Callable[[CallbackType, ItemEvent], UnsubscribeType] +@callback +def async_device_available_fn(controller: UniFiController, obj_id: str) -> bool: + """Check if device is available.""" + if "_" in obj_id: # Sub device (outlet or port) + obj_id = obj_id.partition("_")[0] + + device = controller.api.devices[obj_id] + return controller.available and not device.disabled + + +@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 (outlet or port) + 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), + ) + + @dataclass class UnifiDescription(Generic[HandlerT, DataT]): """Validate and load entities from different UniFi handlers.""" @@ -106,6 +144,15 @@ class UnifiEntity(Entity, Generic[HandlerT, DataT]): ) ) + # Subscribe to events if defined + 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, + ) + ) + @callback def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: """Update the entity state.""" @@ -157,3 +204,11 @@ class UnifiEntity(Entity, Generic[HandlerT, DataT]): Perform additional actions updating platform entity child class state. """ + + @callback + def async_event_callback(self, event: Event) -> None: + """Update entity state based on subscribed event. + + Perform additional action updating platform entity child class state. + """ + raise NotImplementedError() diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 44007e4c1a8..8186c708698 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -9,20 +9,14 @@ from __future__ import annotations import asyncio from collections.abc import Callable, Coroutine from dataclasses import dataclass -from typing import Any, Generic, TypeVar, Union +from typing import Any, Generic import aiounifi -from aiounifi.interfaces.api_handlers import ( - APIHandler, - CallbackType, - ItemEvent, - UnsubscribeType, -) +from aiounifi.interfaces.api_handlers import ItemEvent 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, @@ -42,32 +36,35 @@ from homeassistant.components.switch 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, DeviceEntryType, ) -from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import DeviceInfo, EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN from .controller import UniFiController +from .entity import ( + DataT, + HandlerT, + SubscriptionT, + UnifiEntity, + UnifiEntityDescription, + async_device_available_fn, + async_device_device_info_fn, +) CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED) CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED) -_DataT = TypeVar("_DataT", bound=Union[APIItem, Outlet, Port]) -_HandlerT = TypeVar("_HandlerT", bound=Union[APIHandler, Outlets, Ports]) - -Subscription = Callable[[CallbackType, ItemEvent], UnsubscribeType] - @callback def async_dpi_group_is_on_fn( - api: aiounifi.Controller, dpi_group: DPIRestrictionGroup + controller: UniFiController, dpi_group: DPIRestrictionGroup ) -> bool: """Calculate if all apps are enabled.""" + api = controller.api return all( api.dpi_apps[app_id].enabled for app_id in dpi_group.dpiapp_ids or [] @@ -75,14 +72,6 @@ def async_dpi_group_is_on_fn( ) -@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.""" @@ -94,23 +83,6 @@ def async_client_device_info_fn(api: aiounifi.Controller, obj_id: str) -> Device ) -@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.""" @@ -163,35 +135,27 @@ async def async_poe_port_control_fn( @dataclass -class UnifiEntityLoader(Generic[_HandlerT, _DataT]): +class UnifiSwitchEntityDescriptionMixin(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, 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, _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] + is_on_fn: Callable[[UniFiController, DataT], bool] @dataclass -class UnifiEntityDescription( - SwitchEntityDescription, UnifiEntityLoader[_HandlerT, _DataT] +class UnifiSwitchEntityDescription( + SwitchEntityDescription, + UnifiEntityDescription[HandlerT, DataT], + UnifiSwitchEntityDescriptionMixin[HandlerT, DataT], ): """Class describing UniFi switch entity.""" - custom_subscribe: Callable[[aiounifi.Controller], Subscription] | None = None + custom_subscribe: Callable[[aiounifi.Controller], SubscriptionT] | None = None only_event_for_state_change: bool = False -ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( - UnifiEntityDescription[Clients, Client]( +ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( + UnifiSwitchEntityDescription[Clients, Client]( key="Block client", device_class=SwitchDeviceClass.SWITCH, entity_category=EntityCategory.CONFIG, @@ -204,14 +168,14 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( 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, + is_on_fn=lambda controller, client: not client.blocked, name_fn=lambda client: None, object_fn=lambda api, obj_id: api.clients[obj_id], only_event_for_state_change=True, - supported_fn=lambda api, obj_id: True, - unique_id_fn=lambda obj_id: f"block-{obj_id}", + supported_fn=lambda controller, obj_id: True, + unique_id_fn=lambda controller, obj_id: f"block-{obj_id}", ), - UnifiEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup]( + UnifiSwitchEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup]( key="DPI restriction", entity_category=EntityCategory.CONFIG, icon="mdi:network", @@ -226,27 +190,27 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( 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, + supported_fn=lambda c, obj_id: bool(c.api.dpi_groups[obj_id].dpiapp_ids), + unique_id_fn=lambda controller, obj_id: obj_id, ), - UnifiEntityDescription[Outlets, Outlet]( + UnifiSwitchEntityDescription[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, + available_fn=async_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, + is_on_fn=lambda controller, 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]}", + supported_fn=lambda c, obj_id: c.api.outlets[obj_id].has_relay, + unique_id_fn=lambda controller, obj_id: f"{obj_id.split('_', 1)[0]}-outlet-{obj_id.split('_', 1)[1]}", ), - UnifiEntityDescription[Ports, Port]( + UnifiSwitchEntityDescription[Ports, Port]( key="PoE port control", device_class=SwitchDeviceClass.OUTLET, entity_category=EntityCategory.CONFIG, @@ -255,16 +219,16 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( icon="mdi:ethernet", allowed_fn=lambda controller, obj_id: True, api_handler_fn=lambda api: api.ports, - available_fn=async_sub_device_available_fn, + available_fn=async_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", + is_on_fn=lambda controller, 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]}", + supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe, + unique_id_fn=lambda controller, obj_id: f"{obj_id.split('_', 1)[0]}-poe-{obj_id.split('_', 1)[1]}", ), ) @@ -285,62 +249,24 @@ async def async_setup_entry( client = controller.api.clients_all[mac] controller.api.clients.process_raw([client.raw]) - @callback - def async_load_entities(description: UnifiEntityDescription) -> None: - """Load and subscribe to UniFi devices.""" - entities: list[SwitchEntity] = [] - api_handler = description.api_handler_fn(controller.api) - - @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 - - entity = UnifiSwitchEntity(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) + controller.register_platform_add_entities( + UnifiSwitchEntity, ENTITY_DESCRIPTIONS, async_add_entities + ) -class UnifiSwitchEntity(SwitchEntity, Generic[_HandlerT, _DataT]): +class UnifiSwitchEntity(UnifiEntity[HandlerT, DataT], SwitchEntity): """Base representation of a UniFi switch.""" - entity_description: UnifiEntityDescription[_HandlerT, _DataT] - _attr_should_poll = False + entity_description: UnifiSwitchEntityDescription[HandlerT, DataT] + only_event_for_state_change = False - 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(obj_id) - - obj = description.object_fn(self.controller.api, obj_id) - self._attr_is_on = description.is_on_fn(controller.api, obj) - self._attr_name = description.name_fn(obj) + @callback + def async_initiate_state(self) -> None: + """Initiate entity state.""" + self.async_update_state(ItemEvent.ADDED, self._obj_id) + self.only_event_for_state_change = ( + self.entity_description.only_event_for_state_change + ) async def async_turn_on(self, **kwargs: Any) -> None: """Turn on switch.""" @@ -354,72 +280,19 @@ class UnifiSwitchEntity(SwitchEntity, Generic[_HandlerT, _DataT]): self.controller.api, self._obj_id, False ) - 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_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 switch state.""" - if event == ItemEvent.DELETED and obj_id == self._obj_id: - self.hass.async_create_task(self.remove_item({self._obj_id})) + def async_update_state(self, event: ItemEvent, obj_id: str) -> None: + """Update entity state. + + Update attr_is_on. + """ + if self.only_event_for_state_change: return description = self.entity_description - if not description.supported_fn(self.controller.api, self._obj_id): - self.hass.async_create_task(self.remove_item({self._obj_id})) - return - - if not description.only_event_for_state_change: - obj = description.object_fn(self.controller.api, self._obj_id) - 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 - def async_signal_reachable_callback(self) -> None: - """Call when controller connection state change.""" - self.async_signalling_callback(ItemEvent.ADDED, self._obj_id) + obj = description.object_fn(self.controller.api, self._obj_id) + if (is_on := description.is_on_fn(self.controller, obj)) != self.is_on: + self._attr_is_on = is_on @callback def async_event_callback(self, event: Event) -> None: @@ -436,17 +309,13 @@ class UnifiSwitchEntity(SwitchEntity, Generic[_HandlerT, _DataT]): self._attr_available = description.available_fn(self.controller, self._obj_id) self.async_write_ha_state() - async def 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 async_added_to_hass(self) -> None: + """Register callbacks.""" + await super().async_added_to_hass() - 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) + if self.entity_description.custom_subscribe is not None: + self.async_on_remove( + self.entity_description.custom_subscribe(self.controller.api)( + self.async_signalling_callback, ItemEvent.CHANGED + ), + ) diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index 36a94b73da2..ea02b144a2f 100644 --- a/homeassistant/components/unifi/update.py +++ b/homeassistant/components/unifi/update.py @@ -4,7 +4,7 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from dataclasses import dataclass import logging -from typing import TYPE_CHECKING, Any, Generic +from typing import TYPE_CHECKING, Any, Generic, TypeVar import aiounifi from aiounifi.interfaces.api_handlers import ItemEvent @@ -19,24 +19,23 @@ from homeassistant.components.update import ( ) from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback -from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC -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 +from .const import DOMAIN as UNIFI_DOMAIN +from .entity import ( + UnifiEntity, + UnifiEntityDescription, + async_device_available_fn, + async_device_device_info_fn, +) if TYPE_CHECKING: from .controller import UniFiController LOGGER = logging.getLogger(__name__) - -@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 +_DataT = TypeVar("_DataT", bound=Device) +_HandlerT = TypeVar("_HandlerT", bound=Devices) async def async_device_control_fn(api: aiounifi.Controller, obj_id: str) -> None: @@ -44,33 +43,19 @@ async def async_device_control_fn(api: aiounifi.Controller, obj_id: str) -> None 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]): +class UnifiUpdateEntityDescriptionMixin(Generic[_HandlerT, _DataT]): """Validate and load entities from different UniFi handlers.""" control_fn: Callable[[aiounifi.Controller, str], Coroutine[Any, Any, None]] - state_fn: Callable[[aiounifi.Controller, DataT], bool] + state_fn: Callable[[aiounifi.Controller, _DataT], bool] @dataclass class UnifiUpdateEntityDescription( UpdateEntityDescription, - UnifiEntityDescription[HandlerT, DataT], - UnifiEntityLoader[HandlerT, DataT], + UnifiEntityDescription[_HandlerT, _DataT], + UnifiUpdateEntityDescriptionMixin[_HandlerT, _DataT], ): """Class describing UniFi update entity.""" @@ -108,10 +93,10 @@ async def async_setup_entry( ) -class UnifiDeviceUpdateEntity(UnifiEntity[HandlerT, DataT], UpdateEntity): +class UnifiDeviceUpdateEntity(UnifiEntity[_HandlerT, _DataT], UpdateEntity): """Representation of a UniFi device update entity.""" - entity_description: UnifiUpdateEntityDescription[HandlerT, DataT] + entity_description: UnifiUpdateEntityDescription[_HandlerT, _DataT] @callback def async_initiate_state(self) -> None: