UniFi library controls add/update signalling (#89525)

* Library controls add/update signalling

* Remove add/remove signalling

* Remove unifi_entity_base and unifi_client to make mypy pass
This commit is contained in:
Robert Svensson 2023-03-11 20:14:39 +01:00 committed by GitHub
parent 7487a004fd
commit 8564768d9e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 8 additions and 255 deletions

View File

@ -10,7 +10,7 @@ from typing import Any
from aiohttp import CookieJar
import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.messages import DATA_CLIENT_REMOVED, DATA_EVENT
from aiounifi.interfaces.messages import DATA_EVENT
from aiounifi.models.event import EventKey
from aiounifi.websocket import WebsocketSignal, WebsocketState
import async_timeout
@ -73,17 +73,6 @@ from .errors import AuthenticationRequired, CannotConnect
RETRY_TIMER = 15
CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1)
CLIENT_CONNECTED = (
EventKey.WIRED_CLIENT_CONNECTED,
EventKey.WIRELESS_CLIENT_CONNECTED,
EventKey.WIRELESS_GUEST_CONNECTED,
)
DEVICE_CONNECTED = (
EventKey.ACCESS_POINT_CONNECTED,
EventKey.GATEWAY_CONNECTED,
EventKey.SWITCH_CONNECTED,
)
class UniFiController:
"""Manages a single UniFi Network instance."""
@ -258,55 +247,20 @@ class UniFiController:
else:
LOGGER.info("Connected to UniFi Network")
elif signal == WebsocketSignal.DATA and data:
if DATA_EVENT in data:
clients_connected = set()
devices_connected = set()
wireless_clients_connected = False
for event in data[DATA_EVENT]:
if event.key in CLIENT_CONNECTED:
clients_connected.add(event.mac)
if not wireless_clients_connected and event.key in (
EventKey.WIRELESS_CLIENT_CONNECTED,
EventKey.WIRELESS_GUEST_CONNECTED,
):
wireless_clients_connected = True
elif event.key in DEVICE_CONNECTED:
devices_connected.add(event.mac)
if wireless_clients_connected:
elif signal == WebsocketSignal.DATA and DATA_EVENT in data:
for event in data[DATA_EVENT]:
if event.key in (
EventKey.WIRELESS_CLIENT_CONNECTED,
EventKey.WIRELESS_GUEST_CONNECTED,
):
self.update_wireless_clients()
if clients_connected or devices_connected:
async_dispatcher_send(
self.hass,
self.signal_update,
clients_connected,
devices_connected,
)
elif DATA_CLIENT_REMOVED in data:
async_dispatcher_send(
self.hass, self.signal_remove, data[DATA_CLIENT_REMOVED]
)
break
@property
def signal_reachable(self) -> str:
"""Integration specific event to signal a change in connection status."""
return f"unifi-reachable-{self.config_entry.entry_id}"
@property
def signal_update(self) -> str:
"""Event specific per UniFi entry to signal new data."""
return f"unifi-update-{self.config_entry.entry_id}"
@property
def signal_remove(self) -> str:
"""Event specific per UniFi entry to signal removal of entities."""
return f"unifi-remove-{self.config_entry.entry_id}"
@property
def signal_options_update(self) -> str:
"""Event specific per UniFi entry to signal new options."""

View File

@ -1,58 +0,0 @@
"""Base class for UniFi clients."""
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.entity import DeviceInfo
from .unifi_entity_base import UniFiBase
class UniFiClientBase(UniFiBase):
"""Base class for UniFi clients (without device info)."""
def __init__(self, client, controller) -> None:
"""Set up client."""
super().__init__(client, controller)
self._is_wired = client.mac not in controller.wireless_clients
self.client = self._item
@property
def is_wired(self):
"""Return if the client is wired.
Allows disabling logic to keep track of clients affected by UniFi wired bug marking wireless devices as wired. This is useful when running a network not only containing UniFi APs.
"""
if self._is_wired and self.client.mac in self.controller.wireless_clients:
self._is_wired = False
if self.controller.option_ignore_wired_bug:
return self.client.is_wired
return self._is_wired
@property
def unique_id(self):
"""Return a unique identifier for this switch."""
return f"{self.TYPE}-{self.client.mac}"
@property
def name(self) -> str:
"""Return the name of the client."""
return self.client.name or self.client.hostname
@property
def available(self) -> bool:
"""Return if controller is available."""
return self.controller.available
class UniFiClient(UniFiClientBase):
"""Base class for UniFi clients (with device info)."""
@property
def device_info(self) -> DeviceInfo:
"""Return a client description for device registry."""
return DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, self.client.mac)},
default_manufacturer=self.client.oui,
default_name=self.client.name or self.client.hostname,
)

View File

@ -1,97 +0,0 @@
"""Base class for UniFi Network entities."""
from __future__ import annotations
from collections.abc import Callable
import logging
from typing import TYPE_CHECKING, Any
from homeassistant.core import callback
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import Entity
if TYPE_CHECKING:
from .controller import UniFiController
_LOGGER = logging.getLogger(__name__)
class UniFiBase(Entity):
"""UniFi entity base class."""
_attr_should_poll = False
DOMAIN = ""
TYPE = ""
def __init__(self, item, controller: UniFiController) -> None:
"""Set up UniFi Network entity base.
Register mac to controller entities to cover disabled entities.
"""
self._item = item
self.controller = controller
self.controller.entities[self.DOMAIN][self.TYPE].add(self.key)
@property
def key(self) -> Any:
"""Return item key."""
return self._item.mac
async def async_added_to_hass(self) -> None:
"""Entity created."""
_LOGGER.debug(
"New %s entity %s (%s)",
self.TYPE,
self.entity_id,
self.key,
)
signals: tuple[tuple[str, Callable[..., Any]], ...] = (
(self.controller.signal_reachable, self.async_signal_reachable_callback),
(self.controller.signal_options_update, self.options_updated),
(self.controller.signal_remove, self.remove_item),
)
for signal, method in signals:
self.async_on_remove(async_dispatcher_connect(self.hass, signal, method))
self._item.register_callback(self.async_update_callback)
async def async_will_remove_from_hass(self) -> None:
"""Disconnect object when removed."""
_LOGGER.debug(
"Removing %s entity %s (%s)",
self.TYPE,
self.entity_id,
self.key,
)
self._item.remove_callback(self.async_update_callback)
self.controller.entities[self.DOMAIN][self.TYPE].remove(self.key)
@callback
def async_signal_reachable_callback(self) -> None:
"""Call when controller connection state change."""
self.async_update_callback()
@callback
def async_update_callback(self) -> None:
"""Update the entity's state."""
_LOGGER.debug(
"Updating %s entity %s (%s)",
self.TYPE,
self.entity_id,
self.key,
)
self.async_write_ha_state()
async def options_updated(self) -> None:
"""Config entry options are updated, remove entity if option is disabled."""
raise NotImplementedError
async def remove_item(self, keys: set) -> None:
"""Remove entity if key is part of set."""
if self.key 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)

View File

@ -244,8 +244,6 @@ async def test_controller_setup(
assert controller.mac is None
assert controller.signal_reachable == "unifi-reachable-1"
assert controller.signal_update == "unifi-update-1"
assert controller.signal_remove == "unifi-remove-1"
assert controller.signal_options_update == "unifi-options-1"
assert controller.signal_heartbeat_missed == "unifi-heartbeat-missed"

View File

@ -14,13 +14,11 @@ from homeassistant.components.unifi.const import (
CONF_ALLOW_UPTIME_SENSORS,
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
DOMAIN as UNIFI_DOMAIN,
)
from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE, EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.entity_registry import RegistryEntryDisabler
import homeassistant.util.dt as dt_util
@ -198,24 +196,6 @@ async def test_bandwidth_sensors(
assert hass.states.get("sensor.wired_client_rx")
assert hass.states.get("sensor.wired_client_tx")
# Try to add the sensors again, using a signal
clients_connected = {wired_client["mac"], wireless_client["mac"]}
devices_connected = set()
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
async_dispatcher_send(
hass,
controller.signal_update,
clients_connected,
devices_connected,
)
await hass.async_block_till_done()
assert len(hass.states.async_all()) == 5
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4
@pytest.mark.parametrize(
("initial_uptime", "event_uptime", "new_uptime"),
@ -311,24 +291,6 @@ async def test_uptime_sensors(
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
assert hass.states.get("sensor.client1_uptime")
# Try to add the sensors again, using a signal
clients_connected = {uptime_client["mac"]}
devices_connected = set()
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
async_dispatcher_send(
hass,
controller.signal_update,
clients_connected,
devices_connected,
)
await hass.async_block_till_done()
assert len(hass.states.async_all()) == 2
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
async def test_remove_sensors(
hass: HomeAssistant,

View File

@ -30,7 +30,6 @@ from homeassistant.const import (
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.entity_registry import RegistryEntryDisabler
from homeassistant.util import dt
@ -719,11 +718,6 @@ async def test_switches(
assert aioclient_mock.call_count == 14
assert aioclient_mock.mock_calls[13][2] == {"enabled": True}
# Make sure no duplicates arise on generic signal update
async_dispatcher_send(hass, controller.signal_update)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3
async def test_remove_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket