From 3444d2af1aff29a09b86b7046d71cb8ba8a1c3b3 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 8 Nov 2022 07:38:31 +0100 Subject: [PATCH] UniFi switch entity description (#81680) * Consolidate switch entities to one class * Move turn on/off into UnifiSwitchEntity * Add event subscription Remove storing entity for everything but legacy poe switch * Only one entity class * Improve generics naming * Rename loader to description * Improve control_fn naming * Move wrongfully placed method that should only react to dpi apps being emptied * Improve different methods * Minor renaming and sorting * Mark callbacks properly --- homeassistant/components/unifi/diagnostics.py | 1 - homeassistant/components/unifi/switch.py | 722 ++++++++---------- tests/components/unifi/test_diagnostics.py | 30 - 3 files changed, 308 insertions(+), 445 deletions(-) diff --git a/homeassistant/components/unifi/diagnostics.py b/homeassistant/components/unifi/diagnostics.py index b35fd520ab0..495613f3b81 100644 --- a/homeassistant/components/unifi/diagnostics.py +++ b/homeassistant/components/unifi/diagnostics.py @@ -95,7 +95,6 @@ async def async_get_config_entry_diagnostics( async_replace_dict_data(config_entry.as_dict(), macs_to_redact), REDACT_CONFIG ) diag["site_role"] = controller.site_role - diag["entities"] = async_replace_dict_data(controller.entities, macs_to_redact) diag["clients"] = { macs_to_redact[k]: async_redact_data( async_replace_dict_data(v.raw, macs_to_redact), REDACT_CLIENTS diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 65d0041187e..e63d5548ebc 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -7,24 +7,33 @@ Support for controlling deep packet inspection (DPI) restriction groups. from __future__ import annotations import asyncio -from collections.abc import Callable +from collections.abc import Callable, Coroutine from dataclasses import dataclass from typing import Any, Generic, TypeVar -from aiounifi.interfaces.api_handlers import ItemEvent +import aiounifi +from aiounifi.interfaces.api_handlers import 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.client import ClientBlockRequest +from aiounifi.models.client import Client, ClientBlockRequest from aiounifi.models.device import ( DeviceSetOutletRelayRequest, DeviceSetPoePortModeRequest, ) from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest +from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup from aiounifi.models.event import Event, EventKey +from aiounifi.models.outlet import Outlet +from aiounifi.models.port import Port -from homeassistant.components.switch import DOMAIN, SwitchDeviceClass, SwitchEntity +from homeassistant.components.switch import ( + DOMAIN, + SwitchDeviceClass, + SwitchEntity, + SwitchEntityDescription, +) from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import entity_registry as er @@ -37,33 +46,219 @@ from homeassistant.helpers.entity import DeviceInfo, EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.restore_state import RestoreEntity -from .const import ( - ATTR_MANUFACTURER, - BLOCK_SWITCH, - DOMAIN as UNIFI_DOMAIN, - DPI_SWITCH, - OUTLET_SWITCH, - POE_SWITCH, -) +from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN, POE_SWITCH from .controller import UniFiController from .unifi_client import UniFiClient CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED) CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED) -T = TypeVar("T") +Data = TypeVar("Data") +Handler = TypeVar("Handler") + +Subscription = Callable[[CallbackType, ItemEvent], UnsubscribeType] + + +@callback +def async_dpi_group_is_on_fn( + api: aiounifi.Controller, dpi_group: DPIRestrictionGroup +) -> bool: + """Calculate if all apps are enabled.""" + return all( + api.dpi_apps[app_id].enabled + for app_id in dpi_group.dpiapp_ids or [] + if app_id in api.dpi_apps + ) + + +@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.""" + client = api.clients[obj_id] + return DeviceInfo( + connections={(CONNECTION_NETWORK_MAC, obj_id)}, + default_manufacturer=client.oui, + default_name=client.name or client.hostname, + ) + + +@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.""" + return DeviceInfo( + entry_type=DeviceEntryType.SERVICE, + identifiers={(DOMAIN, f"unifi_controller_{obj_id}")}, + manufacturer=ATTR_MANUFACTURER, + model="UniFi Network", + name="UniFi Network", + ) + + +async def async_block_client_control_fn( + api: aiounifi.Controller, obj_id: str, target: bool +) -> None: + """Control network access of client.""" + await api.request(ClientBlockRequest.create(obj_id, not target)) + + +async def async_dpi_group_control_fn( + api: aiounifi.Controller, obj_id: str, target: bool +) -> None: + """Enable or disable DPI group.""" + dpi_group = api.dpi_groups[obj_id] + await asyncio.gather( + *[ + api.request(DPIRestrictionAppEnableRequest.create(app_id, target)) + for app_id in dpi_group.dpiapp_ids or [] + ] + ) + + +async def async_outlet_control_fn( + api: aiounifi.Controller, obj_id: str, target: bool +) -> None: + """Control outlet relay.""" + mac, _, index = obj_id.partition("_") + device = api.devices[mac] + await api.request(DeviceSetOutletRelayRequest.create(device, int(index), target)) + + +async def async_poe_port_control_fn( + api: aiounifi.Controller, obj_id: str, target: bool +) -> None: + """Control poe state.""" + mac, _, index = obj_id.partition("_") + device = api.devices[mac] + state = "auto" if target else "off" + await api.request(DeviceSetPoePortModeRequest.create(device, int(index), state)) @dataclass -class UnifiEntityLoader(Generic[T]): +class UnifiEntityLoader(Generic[Handler, Data]): """Validate and load entities from different UniFi handlers.""" allowed_fn: Callable[[UniFiController, str], bool] - entity_cls: type[UnifiBlockClientSwitch] | type[UnifiDPIRestrictionSwitch] | type[ - UnifiOutletSwitch - ] | type[UnifiPoePortSwitch] | type[UnifiDPIRestrictionSwitch] - handler_fn: Callable[[UniFiController], T] - supported_fn: Callable[[T, str], bool | None] + api_handler_fn: Callable[[aiounifi.Controller], Handler] + 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] + supported_fn: Callable[[aiounifi.Controller, str], bool | None] + unique_id_fn: Callable[[str], str] + + +@dataclass +class UnifiEntityDescription(SwitchEntityDescription, UnifiEntityLoader[Handler, Data]): + """Class describing UniFi switch entity.""" + + custom_subscribe: Callable[[aiounifi.Controller], Subscription] | None = None + + +ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = ( + UnifiEntityDescription[Clients, Client]( + key="Block client", + device_class=SwitchDeviceClass.SWITCH, + entity_category=EntityCategory.CONFIG, + has_entity_name=True, + icon="mdi:ethernet", + allowed_fn=lambda controller, obj_id: obj_id in controller.option_block_clients, + api_handler_fn=lambda api: api.clients, + available_fn=lambda controller, obj_id: controller.available, + control_fn=async_block_client_control_fn, + 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, + name_fn=lambda client: None, + object_fn=lambda api, obj_id: api.clients[obj_id], + supported_fn=lambda api, obj_id: True, + unique_id_fn=lambda obj_id: f"block-{obj_id}", + ), + UnifiEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup]( + key="DPI restriction", + entity_category=EntityCategory.CONFIG, + icon="mdi:network", + allowed_fn=lambda controller, obj_id: controller.option_dpi_restrictions, + api_handler_fn=lambda api: api.dpi_groups, + available_fn=lambda controller, obj_id: controller.available, + control_fn=async_dpi_group_control_fn, + custom_subscribe=lambda api: api.dpi_apps.subscribe, + device_info_fn=async_dpi_group_device_info_fn, + event_is_on=None, + event_to_subscribe=None, + 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, + ), + UnifiEntityDescription[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, + 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, + 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]}", + ), + UnifiEntityDescription[Ports, Port]( + key="PoE port control", + device_class=SwitchDeviceClass.OUTLET, + entity_category=EntityCategory.CONFIG, + has_entity_name=True, + entity_registry_enabled_default=False, + icon="mdi:ethernet", + allowed_fn=lambda controller, obj_id: True, + api_handler_fn=lambda api: api.ports, + available_fn=async_sub_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", + 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]}", + ), +) async def async_setup_entry( @@ -71,17 +266,9 @@ async def async_setup_entry( config_entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: - """Set up switches for UniFi Network integration. - - Switches are controlling network access and switch ports with POE. - """ + """Set up switches for UniFi Network integration.""" controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] - controller.entities[DOMAIN] = { - BLOCK_SWITCH: set(), - POE_SWITCH: set(), - DPI_SWITCH: set(), - OUTLET_SWITCH: set(), - } + controller.entities[DOMAIN] = {POE_SWITCH: set()} if controller.site_role != "admin": return @@ -125,20 +312,20 @@ async def async_setup_entry( known_poe_clients.clear() @callback - def async_load_entities(loader: UnifiEntityLoader) -> None: + def async_load_entities(description: UnifiEntityDescription) -> None: """Load and subscribe to UniFi devices.""" entities: list[SwitchEntity] = [] - api_handler = loader.handler_fn(controller) + api_handler = description.api_handler_fn(controller.api) @callback def async_create_entity(event: ItemEvent, obj_id: str) -> None: """Create UniFi entity.""" - if not loader.allowed_fn(controller, obj_id) or not loader.supported_fn( - api_handler, obj_id - ): + if not description.allowed_fn( + controller, obj_id + ) or not description.supported_fn(controller.api, obj_id): return - entity = loader.entity_cls(obj_id, controller) + entity = UnifiSwitchEntity(obj_id, controller, description) if event == ItemEvent.ADDED: async_add_entities([entity]) return @@ -150,8 +337,8 @@ async def async_setup_entry( api_handler.subscribe(async_create_entity, ItemEvent.ADDED) - for unifi_loader in UNIFI_LOADERS: - async_load_entities(unifi_loader) + for description in ENTITY_DESCRIPTIONS: + async_load_entities(description) @callback @@ -301,51 +488,52 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity): await self.remove_item({self.client.mac}) -class UnifiBlockClientSwitch(SwitchEntity): - """Representation of a blockable client.""" +class UnifiSwitchEntity(SwitchEntity): + """Base representation of a UniFi switch.""" - _attr_device_class = SwitchDeviceClass.SWITCH - _attr_entity_category = EntityCategory.CONFIG - _attr_has_entity_name = True - _attr_icon = "mdi:ethernet" + entity_description: UnifiEntityDescription _attr_should_poll = False - def __init__(self, obj_id: str, controller: UniFiController) -> None: - """Set up block switch.""" - controller.entities[DOMAIN][BLOCK_SWITCH].add(obj_id) + def __init__( + self, + obj_id: str, + controller: UniFiController, + description: UnifiEntityDescription, + ) -> None: + """Set up UniFi switch entity.""" self._obj_id = obj_id self.controller = controller + self.entity_description = description self._removed = False - client = controller.api.clients[obj_id] - self._attr_available = controller.available - self._attr_is_on = not client.blocked - self._attr_unique_id = f"{BLOCK_SWITCH}-{obj_id}" - self._attr_device_info = DeviceInfo( - connections={(CONNECTION_NETWORK_MAC, obj_id)}, - default_manufacturer=client.oui, - default_name=client.name or client.hostname, + 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: + """Turn on switch.""" + await self.entity_description.control_fn( + self.controller.api, self._obj_id, True + ) + + async def async_turn_off(self, **kwargs: Any) -> None: + """Turn off switch.""" + await self.entity_description.control_fn( + self.controller.api, self._obj_id, False ) 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.clients.subscribe(self.async_signalling_callback) - ) - self.async_on_remove( - self.controller.api.events.subscribe( - self.async_event_callback, CLIENT_BLOCKED + CLIENT_UNBLOCKED - ) - ) - self.async_on_remove( - async_dispatcher_connect( - self.hass, self.controller.signal_remove, self.remove_item - ) - ) - self.async_on_remove( - async_dispatcher_connect( - self.hass, self.controller.signal_options_update, self.options_updated + handler.subscribe( + self.async_signalling_callback, ) ) self.async_on_remove( @@ -355,31 +543,49 @@ class UnifiBlockClientSwitch(SwitchEntity): self.async_signal_reachable_callback, ) ) - - async def async_will_remove_from_hass(self) -> None: - """Disconnect object when removed.""" - self.controller.entities[DOMAIN][BLOCK_SWITCH].remove(self._obj_id) + 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 clients state.""" - if event == ItemEvent.DELETED: + """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 - client = self.controller.api.clients[self._obj_id] - self._attr_is_on = not client.blocked - self._attr_available = self.controller.available - self.async_write_ha_state() - - @callback - def async_event_callback(self, event: Event) -> None: - """Event subscription callback.""" - if event.mac != self._obj_id: + 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 event.key in CLIENT_BLOCKED + CLIENT_UNBLOCKED: - self._attr_is_on = event.key in CLIENT_UNBLOCKED - self._attr_available = self.controller.available + + 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 @@ -387,30 +593,28 @@ class UnifiBlockClientSwitch(SwitchEntity): """Call when controller connection state change.""" self.async_signalling_callback(ItemEvent.ADDED, self._obj_id) - async def async_turn_on(self, **kwargs: Any) -> None: - """Turn on connectivity for client.""" - await self.controller.api.request( - ClientBlockRequest.create(self._obj_id, False) - ) + @callback + def async_event_callback(self, event: Event) -> None: + """Event subscription callback.""" + if event.mac != self._obj_id: + return - async def async_turn_off(self, **kwargs: Any) -> None: - """Turn off connectivity for client.""" - await self.controller.api.request(ClientBlockRequest.create(self._obj_id, True)) + description = self.entity_description + assert isinstance(description.event_to_subscribe, tuple) + assert isinstance(description.event_is_on, tuple) - @property - def icon(self) -> str: - """Return the icon to use in the frontend.""" - if not self.is_on: - return "mdi:network-off" - return "mdi:network" + if event.key in description.event_to_subscribe: + self._attr_is_on = event.key in description.event_is_on + 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 self._obj_id not in self.controller.option_block_clients: + 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 key is part of set.""" + """Remove entity if object ID is part of set.""" if self._obj_id not in keys or self._removed: return self._removed = True @@ -418,313 +622,3 @@ class UnifiBlockClientSwitch(SwitchEntity): er.async_get(self.hass).async_remove(self.entity_id) else: await self.async_remove(force_remove=True) - - -class UnifiDPIRestrictionSwitch(SwitchEntity): - """Representation of a DPI restriction group.""" - - _attr_entity_category = EntityCategory.CONFIG - - def __init__(self, obj_id: str, controller: UniFiController) -> None: - """Set up dpi switch.""" - controller.entities[DOMAIN][DPI_SWITCH].add(obj_id) - self._obj_id = obj_id - self.controller = controller - - dpi_group = controller.api.dpi_groups[obj_id] - self._known_app_ids = dpi_group.dpiapp_ids - - self._attr_available = controller.available - self._attr_is_on = self.calculate_enabled() - self._attr_name = dpi_group.name - self._attr_unique_id = dpi_group.id - self._attr_device_info = DeviceInfo( - entry_type=DeviceEntryType.SERVICE, - identifiers={(DOMAIN, f"unifi_controller_{obj_id}")}, - manufacturer=ATTR_MANUFACTURER, - model="UniFi Network", - name="UniFi Network", - ) - - async def async_added_to_hass(self) -> None: - """Register callback to known apps.""" - self.async_on_remove( - self.controller.api.dpi_groups.subscribe(self.async_signalling_callback) - ) - self.async_on_remove( - self.controller.api.dpi_apps.subscribe( - self.async_signalling_callback, ItemEvent.CHANGED - ), - ) - self.async_on_remove( - async_dispatcher_connect( - self.hass, self.controller.signal_remove, self.remove_item - ) - ) - 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_reachable, - self.async_signal_reachable_callback, - ) - ) - - async def async_will_remove_from_hass(self) -> None: - """Disconnect object when removed.""" - self.controller.entities[DOMAIN][DPI_SWITCH].remove(self._obj_id) - - @callback - def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: - """Object has new event.""" - if event == ItemEvent.DELETED: - self.hass.async_create_task(self.remove_item({self._obj_id})) - return - - dpi_group = self.controller.api.dpi_groups[self._obj_id] - if not dpi_group.dpiapp_ids: - self.hass.async_create_task(self.remove_item({self._obj_id})) - return - - self._attr_available = self.controller.available - self._attr_is_on = self.calculate_enabled() - 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) - - @property - def icon(self): - """Return the icon to use in the frontend.""" - if self.is_on: - return "mdi:network" - return "mdi:network-off" - - def calculate_enabled(self) -> bool: - """Calculate if all apps are enabled.""" - dpi_group = self.controller.api.dpi_groups[self._obj_id] - return all( - self.controller.api.dpi_apps[app_id].enabled - for app_id in dpi_group.dpiapp_ids - if app_id in self.controller.api.dpi_apps - ) - - async def async_turn_on(self, **kwargs: Any) -> None: - """Restrict access of apps related to DPI group.""" - dpi_group = self.controller.api.dpi_groups[self._obj_id] - return await asyncio.gather( - *[ - self.controller.api.request( - DPIRestrictionAppEnableRequest.create(app_id, True) - ) - for app_id in dpi_group.dpiapp_ids - ] - ) - - async def async_turn_off(self, **kwargs: Any) -> None: - """Remove restriction of apps related to DPI group.""" - dpi_group = self.controller.api.dpi_groups[self._obj_id] - return await asyncio.gather( - *[ - self.controller.api.request( - DPIRestrictionAppEnableRequest.create(app_id, False) - ) - for app_id in dpi_group.dpiapp_ids - ] - ) - - async def options_updated(self) -> None: - """Config entry options are updated, remove entity if option is disabled.""" - if not self.controller.option_dpi_restrictions: - await self.remove_item({self._attr_unique_id}) - - async def remove_item(self, keys: set) -> None: - """Remove entity if key is part of set.""" - if self._attr_unique_id not in keys: - return - - if self.registry_entry: - er.async_get(self.hass).async_remove(self.entity_id) - else: - await self.async_remove(force_remove=True) - - -class UnifiOutletSwitch(SwitchEntity): - """Representation of a outlet relay.""" - - _attr_device_class = SwitchDeviceClass.OUTLET - _attr_has_entity_name = True - _attr_should_poll = False - - def __init__(self, obj_id: str, controller: UniFiController) -> None: - """Set up UniFi Network entity base.""" - self._device_mac, index = obj_id.split("_", 1) - self._index = int(index) - self._obj_id = obj_id - self.controller = controller - - outlet = self.controller.api.outlets[self._obj_id] - self._attr_name = outlet.name - self._attr_is_on = outlet.relay_state - self._attr_unique_id = f"{self._device_mac}-outlet-{index}" - - device = self.controller.api.devices[self._device_mac] - self._attr_available = controller.available and not device.disabled - self._attr_device_info = DeviceInfo( - connections={(CONNECTION_NETWORK_MAC, device.mac)}, - manufacturer=ATTR_MANUFACTURER, - model=device.model, - name=device.name or None, - sw_version=device.version, - hw_version=device.board_revision, - ) - - async def async_added_to_hass(self) -> None: - """Entity created.""" - self.async_on_remove( - self.controller.api.outlets.subscribe(self.async_signalling_callback) - ) - self.async_on_remove( - async_dispatcher_connect( - self.hass, - self.controller.signal_reachable, - self.async_signal_reachable_callback, - ) - ) - - @callback - def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: - """Object has new event.""" - device = self.controller.api.devices[self._device_mac] - outlet = self.controller.api.outlets[self._obj_id] - self._attr_available = self.controller.available and not device.disabled - self._attr_is_on = outlet.relay_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 async_turn_on(self, **kwargs: Any) -> None: - """Enable outlet relay.""" - device = self.controller.api.devices[self._device_mac] - await self.controller.api.request( - DeviceSetOutletRelayRequest.create(device, self._index, True) - ) - - async def async_turn_off(self, **kwargs: Any) -> None: - """Disable outlet relay.""" - device = self.controller.api.devices[self._device_mac] - await self.controller.api.request( - DeviceSetOutletRelayRequest.create(device, self._index, False) - ) - - -class UnifiPoePortSwitch(SwitchEntity): - """Representation of a Power-over-Ethernet source port on an UniFi device.""" - - _attr_device_class = SwitchDeviceClass.OUTLET - _attr_entity_category = EntityCategory.CONFIG - _attr_entity_registry_enabled_default = False - _attr_has_entity_name = True - _attr_icon = "mdi:ethernet" - _attr_should_poll = False - - def __init__(self, obj_id: str, controller: UniFiController) -> None: - """Set up UniFi Network entity base.""" - self._device_mac, index = obj_id.split("_", 1) - self._index = int(index) - self._obj_id = obj_id - self.controller = controller - - port = self.controller.api.ports[self._obj_id] - self._attr_name = f"{port.name} PoE" - self._attr_is_on = port.poe_mode != "off" - self._attr_unique_id = f"{self._device_mac}-poe-{index}" - - device = self.controller.api.devices[self._device_mac] - self._attr_available = controller.available and not device.disabled - self._attr_device_info = DeviceInfo( - connections={(CONNECTION_NETWORK_MAC, device.mac)}, - manufacturer=ATTR_MANUFACTURER, - model=device.model, - name=device.name or None, - sw_version=device.version, - hw_version=device.board_revision, - ) - - async def async_added_to_hass(self) -> None: - """Entity created.""" - self.async_on_remove( - self.controller.api.ports.subscribe(self.async_signalling_callback) - ) - self.async_on_remove( - async_dispatcher_connect( - self.hass, - self.controller.signal_reachable, - self.async_signal_reachable_callback, - ) - ) - - @callback - def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None: - """Object has new event.""" - device = self.controller.api.devices[self._device_mac] - port = self.controller.api.ports[self._obj_id] - self._attr_available = self.controller.available and not device.disabled - self._attr_is_on = port.poe_mode != "off" - 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_turn_on(self, **kwargs: Any) -> None: - """Enable POE for client.""" - device = self.controller.api.devices[self._device_mac] - await self.controller.api.request( - DeviceSetPoePortModeRequest.create(device, self._index, "auto") - ) - - async def async_turn_off(self, **kwargs: Any) -> None: - """Disable POE for client.""" - device = self.controller.api.devices[self._device_mac] - await self.controller.api.request( - DeviceSetPoePortModeRequest.create(device, self._index, "off") - ) - - -UNIFI_LOADERS: tuple[UnifiEntityLoader, ...] = ( - UnifiEntityLoader[Clients]( - allowed_fn=lambda controller, obj_id: obj_id in controller.option_block_clients, - entity_cls=UnifiBlockClientSwitch, - handler_fn=lambda contrlr: contrlr.api.clients, - supported_fn=lambda handler, obj_id: True, - ), - UnifiEntityLoader[DPIRestrictionGroups]( - allowed_fn=lambda controller, obj_id: controller.option_dpi_restrictions, - entity_cls=UnifiDPIRestrictionSwitch, - handler_fn=lambda controller: controller.api.dpi_groups, - supported_fn=lambda handler, obj_id: bool(handler[obj_id].dpiapp_ids), - ), - UnifiEntityLoader[Outlets]( - allowed_fn=lambda controller, obj_id: True, - entity_cls=UnifiOutletSwitch, - handler_fn=lambda controller: controller.api.outlets, - supported_fn=lambda handler, obj_id: handler[obj_id].has_relay, - ), - UnifiEntityLoader[Ports]( - allowed_fn=lambda controller, obj_id: True, - entity_cls=UnifiPoePortSwitch, - handler_fn=lambda controller: controller.api.ports, - supported_fn=lambda handler, obj_id: handler[obj_id].port_poe, - ), -) diff --git a/tests/components/unifi/test_diagnostics.py b/tests/components/unifi/test_diagnostics.py index 6584b947293..9de0e4b6154 100644 --- a/tests/components/unifi/test_diagnostics.py +++ b/tests/components/unifi/test_diagnostics.py @@ -6,16 +6,6 @@ from homeassistant.components.unifi.const import ( CONF_ALLOW_UPTIME_SENSORS, CONF_BLOCK_CLIENT, ) -from homeassistant.components.unifi.device_tracker import CLIENT_TRACKER, DEVICE_TRACKER -from homeassistant.components.unifi.sensor import RX_SENSOR, TX_SENSOR, UPTIME_SENSOR -from homeassistant.components.unifi.switch import ( - BLOCK_SWITCH, - DPI_SWITCH, - OUTLET_SWITCH, - POE_SWITCH, -) -from homeassistant.components.unifi.update import DEVICE_UPDATE -from homeassistant.const import Platform from .test_controller import setup_unifi_integration @@ -146,26 +136,6 @@ async def test_entry_diagnostics(hass, hass_client, aioclient_mock): "version": 1, }, "site_role": "admin", - "entities": { - str(Platform.DEVICE_TRACKER): { - CLIENT_TRACKER: ["00:00:00:00:00:00"], - DEVICE_TRACKER: ["00:00:00:00:00:01"], - }, - str(Platform.SENSOR): { - RX_SENSOR: ["00:00:00:00:00:00"], - TX_SENSOR: ["00:00:00:00:00:00"], - UPTIME_SENSOR: ["00:00:00:00:00:00"], - }, - str(Platform.SWITCH): { - BLOCK_SWITCH: ["00:00:00:00:00:00"], - DPI_SWITCH: ["5f976f4ae3c58f018ec7dff6"], - POE_SWITCH: ["00:00:00:00:00:00"], - OUTLET_SWITCH: [], - }, - str(Platform.UPDATE): { - DEVICE_UPDATE: ["00:00:00:00:00:01"], - }, - }, "clients": { "00:00:00:00:00:00": { "blocked": False,