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
This commit is contained in:
Robert Svensson
GitHub
parent
c3d4a9cd99
commit
3444d2af1a
3 changed files with 308 additions and 445 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue