From 220ec1906c0854d5c6cc7b38047805fea02363f1 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 6 Jan 2023 12:50:32 +0100 Subject: [PATCH] Make sensor platform use common UniFi entity class (#84818) fixes undefined --- homeassistant/components/unifi/entity.py | 3 +- homeassistant/components/unifi/sensor.py | 231 +++++------------------ 2 files changed, 51 insertions(+), 183 deletions(-) diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index ff6d368bf97..e7dca396fae 100644 --- a/homeassistant/components/unifi/entity.py +++ b/homeassistant/components/unifi/entity.py @@ -190,12 +190,13 @@ class UnifiEntity(Entity, Generic[HandlerT, DataT]): 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. + Defaults to using async_update_state to set initial state. """ + self.async_update_state(ItemEvent.ADDED, self._obj_id) @callback @abstractmethod diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 95167123295..a88b750fdf7 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -8,7 +8,7 @@ from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Generic, TypeVar, Union +from typing import Generic import aiounifi from aiounifi.interfaces.api_handlers import ItemEvent @@ -25,23 +25,26 @@ from homeassistant.components.sensor import ( from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfInformation, UnitOfPower 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, EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback import homeassistant.util.dt as dt_util -from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN +from .const import DOMAIN as UNIFI_DOMAIN from .controller import UniFiController - -_DataT = TypeVar("_DataT", bound=Union[Client, Port]) -_HandlerT = TypeVar("_HandlerT", bound=Union[Clients, Ports]) +from .entity import ( + DataT, + HandlerT, + UnifiEntity, + UnifiEntityDescription, + async_device_available_fn, + async_device_device_info_fn, +) @callback def async_client_rx_value_fn(controller: UniFiController, client: Client) -> float: - """Calculate if all apps are enabled.""" + """Calculate receiving data transfer value.""" if client.mac not in controller.wireless_clients: return client.wired_rx_bytes_r / 1000000 return client.rx_bytes_r / 1000000 @@ -49,7 +52,7 @@ def async_client_rx_value_fn(controller: UniFiController, client: Client) -> flo @callback def async_client_tx_value_fn(controller: UniFiController, client: Client) -> float: - """Calculate if all apps are enabled.""" + """Calculate transmission data transfer value.""" if client.mac not in controller.wireless_clients: return client.wired_tx_bytes_r / 1000000 return client.tx_bytes_r / 1000000 @@ -76,55 +79,24 @@ 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 device.model, - sw_version=device.version, - hw_version=str(device.board_revision), - ) - - -@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 - - @dataclass -class UnifiEntityLoader(Generic[_HandlerT, _DataT]): +class UnifiSensorEntityDescriptionMixin(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] - name_fn: Callable[[_DataT], str | None] - object_fn: Callable[[aiounifi.Controller, str], _DataT] - supported_fn: Callable[[UniFiController, str], bool | None] - unique_id_fn: Callable[[str], str] - value_fn: Callable[[UniFiController, _DataT], datetime | float | str | None] + value_fn: Callable[[UniFiController, DataT], datetime | float | str | None] @dataclass -class UnifiEntityDescription( - SensorEntityDescription, UnifiEntityLoader[_HandlerT, _DataT] +class UnifiSensorEntityDescription( + SensorEntityDescription, + UnifiEntityDescription[HandlerT, DataT], + UnifiSensorEntityDescriptionMixin[HandlerT, DataT], ): """Class describing UniFi sensor entity.""" -ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( - UnifiEntityDescription[Clients, Client]( +ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = ( + UnifiSensorEntityDescription[Clients, Client]( key="Bandwidth sensor RX", entity_category=EntityCategory.DIAGNOSTIC, native_unit_of_measurement=UnitOfInformation.MEGABYTES, @@ -133,13 +105,15 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( api_handler_fn=lambda api: api.clients, available_fn=lambda controller, _: controller.available, device_info_fn=async_client_device_info_fn, + event_is_on=None, + event_to_subscribe=None, name_fn=lambda _: "RX", object_fn=lambda api, obj_id: api.clients[obj_id], supported_fn=lambda controller, _: controller.option_allow_bandwidth_sensors, - unique_id_fn=lambda obj_id: f"rx-{obj_id}", + unique_id_fn=lambda controller, obj_id: f"rx-{obj_id}", value_fn=async_client_rx_value_fn, ), - UnifiEntityDescription[Clients, Client]( + UnifiSensorEntityDescription[Clients, Client]( key="Bandwidth sensor TX", entity_category=EntityCategory.DIAGNOSTIC, native_unit_of_measurement=UnitOfInformation.MEGABYTES, @@ -148,13 +122,15 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( api_handler_fn=lambda api: api.clients, available_fn=lambda controller, _: controller.available, device_info_fn=async_client_device_info_fn, + event_is_on=None, + event_to_subscribe=None, name_fn=lambda _: "TX", object_fn=lambda api, obj_id: api.clients[obj_id], supported_fn=lambda controller, _: controller.option_allow_bandwidth_sensors, - unique_id_fn=lambda obj_id: f"tx-{obj_id}", + unique_id_fn=lambda controller, obj_id: f"tx-{obj_id}", value_fn=async_client_tx_value_fn, ), - UnifiEntityDescription[Ports, Port]( + UnifiSensorEntityDescription[Ports, Port]( key="PoE port power sensor", device_class=SensorDeviceClass.POWER, entity_category=EntityCategory.DIAGNOSTIC, @@ -163,15 +139,17 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( entity_registry_enabled_default=False, 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, device_info_fn=async_device_device_info_fn, + event_is_on=None, + event_to_subscribe=None, name_fn=lambda port: f"{port.name} PoE Power", object_fn=lambda api, obj_id: api.ports[obj_id], supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe, - unique_id_fn=lambda obj_id: f"poe_power-{obj_id}", + unique_id_fn=lambda controller, obj_id: f"poe_power-{obj_id}", value_fn=lambda _, obj: obj.poe_power if obj.poe_mode != "off" else "0", ), - UnifiEntityDescription[Clients, Client]( + UnifiSensorEntityDescription[Clients, Client]( key="Uptime sensor", device_class=SensorDeviceClass.TIMESTAMP, entity_category=EntityCategory.DIAGNOSTIC, @@ -180,10 +158,12 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( api_handler_fn=lambda api: api.clients, available_fn=lambda controller, obj_id: controller.available, device_info_fn=async_client_device_info_fn, + event_is_on=None, + event_to_subscribe=None, name_fn=lambda client: "Uptime", object_fn=lambda api, obj_id: api.clients[obj_id], supported_fn=lambda controller, _: controller.option_allow_uptime_sensors, - unique_id_fn=lambda obj_id: f"uptime-{obj_id}", + unique_id_fn=lambda controller, obj_id: f"uptime-{obj_id}", value_fn=async_client_uptime_value_fn, ), ) @@ -196,136 +176,23 @@ async def async_setup_entry( ) -> None: """Set up sensors for UniFi Network integration.""" controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + controller.register_platform_add_entities( + UnifiSensorEntity, ENTITY_DESCRIPTIONS, async_add_entities + ) + + +class UnifiSensorEntity(UnifiEntity[HandlerT, DataT], SensorEntity): + """Base representation of a UniFi sensor.""" + + entity_description: UnifiSensorEntityDescription[HandlerT, DataT] @callback - def async_load_entities(description: UnifiEntityDescription) -> None: - """Load and subscribe to UniFi devices.""" - entities: list[SensorEntity] = [] - api_handler = description.api_handler_fn(controller.api) + def async_update_state(self, event: ItemEvent, obj_id: str) -> None: + """Update entity state. - @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, obj_id): - return - - entity = UnifiSensorEntity(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 UnifiSensorEntity(SensorEntity, Generic[_HandlerT, _DataT]): - """Base representation of a UniFi switch.""" - - 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 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(controller.api, obj_id) - self._attr_native_value = description.value_fn(controller, obj) - self._attr_name = description.name_fn(obj) - - async def async_added_to_hass(self) -> None: - """Register callbacks.""" + Update native_value. + """ 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, - ) - ) - - @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 - - 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 - - update_state = False obj = description.object_fn(self.controller.api, self._obj_id) if (value := description.value_fn(self.controller, obj)) != self.native_value: self._attr_native_value = value - update_state = True - if ( - available := description.available_fn(self.controller, self._obj_id) - ) != self.available: - self._attr_available = available - update_state = True - if update_state: - 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 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)