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
This commit is contained in:
Robert Svensson 2023-01-03 22:57:44 +01:00 committed by GitHub
parent 799d527fb5
commit 6718b40181
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 147 additions and 238 deletions

View File

@ -4,27 +4,65 @@ from __future__ import annotations
from abc import abstractmethod from abc import abstractmethod
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING, Generic, TypeVar from typing import TYPE_CHECKING, Generic, TypeVar, Union
import aiounifi import aiounifi
from aiounifi.interfaces.api_handlers import CallbackType, ItemEvent, UnsubscribeType from aiounifi.interfaces.api_handlers import (
from aiounifi.interfaces.devices import Devices APIHandler,
from aiounifi.models.device import Device CallbackType,
from aiounifi.models.event import EventKey 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.core import callback
from homeassistant.helpers import entity_registry as er 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.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import DeviceInfo, Entity, EntityDescription from homeassistant.helpers.entity import DeviceInfo, Entity, EntityDescription
from .const import ATTR_MANUFACTURER
if TYPE_CHECKING: if TYPE_CHECKING:
from .controller import UniFiController from .controller import UniFiController
DataT = TypeVar("DataT", bound=Device) DataT = TypeVar("DataT", bound=Union[APIItem, Outlet, Port])
HandlerT = TypeVar("HandlerT", bound=Devices) HandlerT = TypeVar("HandlerT", bound=Union[APIHandler, Outlets, Ports])
SubscriptionT = Callable[[CallbackType, ItemEvent], UnsubscribeType] 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 @dataclass
class UnifiDescription(Generic[HandlerT, DataT]): class UnifiDescription(Generic[HandlerT, DataT]):
"""Validate and load entities from different UniFi handlers.""" """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 @callback
def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
"""Update the entity state.""" """Update the entity state."""
@ -157,3 +204,11 @@ class UnifiEntity(Entity, Generic[HandlerT, DataT]):
Perform additional actions updating platform entity child class state. 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()

View File

@ -9,20 +9,14 @@ from __future__ import annotations
import asyncio import asyncio
from collections.abc import Callable, Coroutine from collections.abc import Callable, Coroutine
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Generic, TypeVar, Union from typing import Any, Generic
import aiounifi import aiounifi
from aiounifi.interfaces.api_handlers import ( from aiounifi.interfaces.api_handlers import ItemEvent
APIHandler,
CallbackType,
ItemEvent,
UnsubscribeType,
)
from aiounifi.interfaces.clients import Clients from aiounifi.interfaces.clients import Clients
from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups
from aiounifi.interfaces.outlets import Outlets from aiounifi.interfaces.outlets import Outlets
from aiounifi.interfaces.ports import Ports from aiounifi.interfaces.ports import Ports
from aiounifi.models.api import APIItem
from aiounifi.models.client import Client, ClientBlockRequest from aiounifi.models.client import Client, ClientBlockRequest
from aiounifi.models.device import ( from aiounifi.models.device import (
DeviceSetOutletRelayRequest, DeviceSetOutletRelayRequest,
@ -42,32 +36,35 @@ from homeassistant.components.switch import (
) )
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.device_registry import ( from homeassistant.helpers.device_registry import (
CONNECTION_NETWORK_MAC, CONNECTION_NETWORK_MAC,
DeviceEntryType, DeviceEntryType,
) )
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import DeviceInfo, EntityCategory from homeassistant.helpers.entity import DeviceInfo, EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
from .controller import UniFiController 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_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED)
CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED) 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 @callback
def async_dpi_group_is_on_fn( def async_dpi_group_is_on_fn(
api: aiounifi.Controller, dpi_group: DPIRestrictionGroup controller: UniFiController, dpi_group: DPIRestrictionGroup
) -> bool: ) -> bool:
"""Calculate if all apps are enabled.""" """Calculate if all apps are enabled."""
api = controller.api
return all( return all(
api.dpi_apps[app_id].enabled api.dpi_apps[app_id].enabled
for app_id in dpi_group.dpiapp_ids or [] 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 @callback
def async_client_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: def async_client_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
"""Create device registry entry for client.""" """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 @callback
def async_dpi_group_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: def async_dpi_group_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
"""Create device registry entry for DPI group.""" """Create device registry entry for DPI group."""
@ -163,35 +135,27 @@ async def async_poe_port_control_fn(
@dataclass @dataclass
class UnifiEntityLoader(Generic[_HandlerT, _DataT]): class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, DataT]):
"""Validate and load entities from different UniFi handlers.""" """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]] control_fn: Callable[[aiounifi.Controller, str, bool], Coroutine[Any, Any, None]]
device_info_fn: Callable[[aiounifi.Controller, str], DeviceInfo] is_on_fn: Callable[[UniFiController, DataT], bool]
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]
@dataclass @dataclass
class UnifiEntityDescription( class UnifiSwitchEntityDescription(
SwitchEntityDescription, UnifiEntityLoader[_HandlerT, _DataT] SwitchEntityDescription,
UnifiEntityDescription[HandlerT, DataT],
UnifiSwitchEntityDescriptionMixin[HandlerT, DataT],
): ):
"""Class describing UniFi switch entity.""" """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 only_event_for_state_change: bool = False
ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
UnifiEntityDescription[Clients, Client]( UnifiSwitchEntityDescription[Clients, Client](
key="Block client", key="Block client",
device_class=SwitchDeviceClass.SWITCH, device_class=SwitchDeviceClass.SWITCH,
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
@ -204,14 +168,14 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = (
device_info_fn=async_client_device_info_fn, device_info_fn=async_client_device_info_fn,
event_is_on=CLIENT_UNBLOCKED, event_is_on=CLIENT_UNBLOCKED,
event_to_subscribe=CLIENT_BLOCKED + 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, name_fn=lambda client: None,
object_fn=lambda api, obj_id: api.clients[obj_id], object_fn=lambda api, obj_id: api.clients[obj_id],
only_event_for_state_change=True, only_event_for_state_change=True,
supported_fn=lambda api, obj_id: True, supported_fn=lambda controller, obj_id: True,
unique_id_fn=lambda obj_id: f"block-{obj_id}", unique_id_fn=lambda controller, obj_id: f"block-{obj_id}",
), ),
UnifiEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup]( UnifiSwitchEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup](
key="DPI restriction", key="DPI restriction",
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
icon="mdi:network", icon="mdi:network",
@ -226,27 +190,27 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = (
is_on_fn=async_dpi_group_is_on_fn, is_on_fn=async_dpi_group_is_on_fn,
name_fn=lambda group: group.name, name_fn=lambda group: group.name,
object_fn=lambda api, obj_id: api.dpi_groups[obj_id], 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), supported_fn=lambda c, obj_id: bool(c.api.dpi_groups[obj_id].dpiapp_ids),
unique_id_fn=lambda obj_id: obj_id, unique_id_fn=lambda controller, obj_id: obj_id,
), ),
UnifiEntityDescription[Outlets, Outlet]( UnifiSwitchEntityDescription[Outlets, Outlet](
key="Outlet control", key="Outlet control",
device_class=SwitchDeviceClass.OUTLET, device_class=SwitchDeviceClass.OUTLET,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda controller, obj_id: True,
api_handler_fn=lambda api: api.outlets, 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, control_fn=async_outlet_control_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=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, name_fn=lambda outlet: outlet.name,
object_fn=lambda api, obj_id: api.outlets[obj_id], object_fn=lambda api, obj_id: api.outlets[obj_id],
supported_fn=lambda api, obj_id: api.outlets[obj_id].has_relay, supported_fn=lambda c, obj_id: c.api.outlets[obj_id].has_relay,
unique_id_fn=lambda obj_id: f"{obj_id.split('_', 1)[0]}-outlet-{obj_id.split('_', 1)[1]}", 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", key="PoE port control",
device_class=SwitchDeviceClass.OUTLET, device_class=SwitchDeviceClass.OUTLET,
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
@ -255,16 +219,16 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = (
icon="mdi:ethernet", icon="mdi:ethernet",
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda controller, obj_id: True,
api_handler_fn=lambda api: api.ports, 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, control_fn=async_poe_port_control_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=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", name_fn=lambda port: f"{port.name} PoE",
object_fn=lambda api, obj_id: api.ports[obj_id], object_fn=lambda api, obj_id: api.ports[obj_id],
supported_fn=lambda api, obj_id: api.ports[obj_id].port_poe, supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe,
unique_id_fn=lambda obj_id: f"{obj_id.split('_', 1)[0]}-poe-{obj_id.split('_', 1)[1]}", 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] client = controller.api.clients_all[mac]
controller.api.clients.process_raw([client.raw]) controller.api.clients.process_raw([client.raw])
@callback controller.register_platform_add_entities(
def async_load_entities(description: UnifiEntityDescription) -> None: UnifiSwitchEntity, ENTITY_DESCRIPTIONS, async_add_entities
"""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)
class UnifiSwitchEntity(SwitchEntity, Generic[_HandlerT, _DataT]): class UnifiSwitchEntity(UnifiEntity[HandlerT, DataT], SwitchEntity):
"""Base representation of a UniFi switch.""" """Base representation of a UniFi switch."""
entity_description: UnifiEntityDescription[_HandlerT, _DataT] entity_description: UnifiSwitchEntityDescription[HandlerT, DataT]
_attr_should_poll = False only_event_for_state_change = False
def __init__( @callback
self, def async_initiate_state(self) -> None:
obj_id: str, """Initiate entity state."""
controller: UniFiController, self.async_update_state(ItemEvent.ADDED, self._obj_id)
description: UnifiEntityDescription[_HandlerT, _DataT], self.only_event_for_state_change = (
) -> None: self.entity_description.only_event_for_state_change
"""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)
async def async_turn_on(self, **kwargs: Any) -> None: async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn on switch.""" """Turn on switch."""
@ -354,72 +280,19 @@ class UnifiSwitchEntity(SwitchEntity, Generic[_HandlerT, _DataT]):
self.controller.api, self._obj_id, False 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 @callback
def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
"""Update the switch state.""" """Update entity state.
if event == ItemEvent.DELETED and obj_id == self._obj_id:
self.hass.async_create_task(self.remove_item({self._obj_id})) Update attr_is_on.
"""
if self.only_event_for_state_change:
return return
description = self.entity_description description = self.entity_description
if not description.supported_fn(self.controller.api, self._obj_id): obj = description.object_fn(self.controller.api, self._obj_id)
self.hass.async_create_task(self.remove_item({self._obj_id})) if (is_on := description.is_on_fn(self.controller, obj)) != self.is_on:
return self._attr_is_on = is_on
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)
@callback @callback
def async_event_callback(self, event: Event) -> None: 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._attr_available = description.available_fn(self.controller, self._obj_id)
self.async_write_ha_state() self.async_write_ha_state()
async def options_updated(self) -> None: async def async_added_to_hass(self) -> None:
"""Config entry options are updated, remove entity if option is disabled.""" """Register callbacks."""
if not self.entity_description.allowed_fn(self.controller, self._obj_id): await super().async_added_to_hass()
await self.remove_item({self._obj_id})
async def remove_item(self, keys: set) -> None: if self.entity_description.custom_subscribe is not None:
"""Remove entity if object ID is part of set.""" self.async_on_remove(
if self._obj_id not in keys or self._removed: self.entity_description.custom_subscribe(self.controller.api)(
return self.async_signalling_callback, ItemEvent.CHANGED
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)

View File

@ -4,7 +4,7 @@ from __future__ import annotations
from collections.abc import Callable, Coroutine from collections.abc import Callable, Coroutine
from dataclasses import dataclass from dataclasses import dataclass
import logging import logging
from typing import TYPE_CHECKING, Any, Generic from typing import TYPE_CHECKING, Any, Generic, TypeVar
import aiounifi import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.api_handlers import ItemEvent
@ -19,24 +19,23 @@ from homeassistant.components.update import (
) )
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback 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 homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN from .const import DOMAIN as UNIFI_DOMAIN
from .entity import DataT, HandlerT, UnifiEntity, UnifiEntityDescription from .entity import (
UnifiEntity,
UnifiEntityDescription,
async_device_available_fn,
async_device_device_info_fn,
)
if TYPE_CHECKING: if TYPE_CHECKING:
from .controller import UniFiController from .controller import UniFiController
LOGGER = logging.getLogger(__name__) LOGGER = logging.getLogger(__name__)
_DataT = TypeVar("_DataT", bound=Device)
@callback _HandlerT = TypeVar("_HandlerT", bound=Devices)
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: 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)) 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 @dataclass
class UnifiEntityLoader(Generic[HandlerT, DataT]): class UnifiUpdateEntityDescriptionMixin(Generic[_HandlerT, _DataT]):
"""Validate and load entities from different UniFi handlers.""" """Validate and load entities from different UniFi handlers."""
control_fn: Callable[[aiounifi.Controller, str], Coroutine[Any, Any, None]] 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 @dataclass
class UnifiUpdateEntityDescription( class UnifiUpdateEntityDescription(
UpdateEntityDescription, UpdateEntityDescription,
UnifiEntityDescription[HandlerT, DataT], UnifiEntityDescription[_HandlerT, _DataT],
UnifiEntityLoader[HandlerT, DataT], UnifiUpdateEntityDescriptionMixin[_HandlerT, _DataT],
): ):
"""Class describing UniFi update entity.""" """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.""" """Representation of a UniFi device update entity."""
entity_description: UnifiUpdateEntityDescription[HandlerT, DataT] entity_description: UnifiUpdateEntityDescription[_HandlerT, _DataT]
@callback @callback
def async_initiate_state(self) -> None: def async_initiate_state(self) -> None: