From 896cd27bea14fdb11d3bde57d7a2902567b85342 Mon Sep 17 00:00:00 2001
From: Kim de Vos <kim.de.vos@vippsmobilepay.com>
Date: Tue, 30 Jul 2024 17:20:56 +0200
Subject: [PATCH] Add sensors for Unifi latency (#116737)
* Add sensors for Unifi latency
* Add needed guard and casting
* Use new types
* Add WAN2 support
* Add literals
* Make ids for WAN and WAN2 unique
* Make methods general
* Update sensor.py
* add more typing
* Simplify usage make_wan_latency_sensors
* Simplify further
* Move latency entity creation to method
* Make method internal
* simplify tests
* Apply feedback
* Reduce boiler copied code and add support function
* Add external method for share logic between
* Remove none
* Add more tests
* Apply feedback and change code to improve code coverage
---
homeassistant/components/unifi/sensor.py | 93 +++++++++-
tests/components/unifi/test_sensor.py | 224 +++++++++++++++++++++++
2 files changed, 315 insertions(+), 2 deletions(-)
diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py
index bb974864f60..d86b72d1b2f 100644
--- a/homeassistant/components/unifi/sensor.py
+++ b/homeassistant/components/unifi/sensor.py
@@ -11,6 +11,7 @@ from dataclasses import dataclass
from datetime import date, datetime, timedelta
from decimal import Decimal
from functools import partial
+from typing import TYPE_CHECKING, Literal
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.clients import Clients
@@ -20,7 +21,7 @@ from aiounifi.interfaces.ports import Ports
from aiounifi.interfaces.wlans import Wlans
from aiounifi.models.api import ApiItemT
from aiounifi.models.client import Client
-from aiounifi.models.device import Device
+from aiounifi.models.device import Device, TypedDeviceUptimeStatsWanMonitor
from aiounifi.models.outlet import Outlet
from aiounifi.models.port import Port
from aiounifi.models.wlan import Wlan
@@ -32,7 +33,13 @@ from homeassistant.components.sensor import (
SensorStateClass,
UnitOfTemperature,
)
-from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfDataRate, UnitOfPower
+from homeassistant.const import (
+ PERCENTAGE,
+ EntityCategory,
+ UnitOfDataRate,
+ UnitOfPower,
+ UnitOfTime,
+)
from homeassistant.core import Event as core_Event, HomeAssistant, callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -192,6 +199,86 @@ def async_device_state_value_fn(hub: UnifiHub, device: Device) -> str:
return DEVICE_STATES[device.state]
+@callback
+def async_device_wan_latency_supported_fn(
+ wan: Literal["WAN", "WAN2"],
+ monitor_target: str,
+ hub: UnifiHub,
+ obj_id: str,
+) -> bool:
+ """Determine if an device have a latency monitor."""
+ if (device := hub.api.devices[obj_id]) and device.uptime_stats:
+ return _device_wan_latency_monitor(wan, monitor_target, device) is not None
+ return False
+
+
+@callback
+def async_device_wan_latency_value_fn(
+ wan: Literal["WAN", "WAN2"],
+ monitor_target: str,
+ hub: UnifiHub,
+ device: Device,
+) -> int | None:
+ """Retrieve the monitor target from WAN monitors."""
+ target = _device_wan_latency_monitor(wan, monitor_target, device)
+
+ if TYPE_CHECKING:
+ # Checked by async_device_wan_latency_supported_fn
+ assert target
+
+ return target.get("latency_average", 0)
+
+
+@callback
+def _device_wan_latency_monitor(
+ wan: Literal["WAN", "WAN2"], monitor_target: str, device: Device
+) -> TypedDeviceUptimeStatsWanMonitor | None:
+ """Return the target of the WAN latency monitor."""
+ if device.uptime_stats and (uptime_stats_wan := device.uptime_stats.get(wan)):
+ for monitor in uptime_stats_wan["monitors"]:
+ if monitor_target in monitor["target"]:
+ return monitor
+ return None
+
+
+def make_wan_latency_sensors() -> tuple[UnifiSensorEntityDescription, ...]:
+ """Create WAN latency sensors from WAN monitor data."""
+
+ def make_wan_latency_entity_description(
+ wan: Literal["WAN", "WAN2"], name: str, monitor_target: str
+ ) -> UnifiSensorEntityDescription:
+ return UnifiSensorEntityDescription[Devices, Device](
+ key=f"{name} {wan} latency",
+ entity_category=EntityCategory.DIAGNOSTIC,
+ native_unit_of_measurement=UnitOfTime.MILLISECONDS,
+ state_class=SensorStateClass.MEASUREMENT,
+ device_class=SensorDeviceClass.DURATION,
+ entity_registry_enabled_default=False,
+ api_handler_fn=lambda api: api.devices,
+ available_fn=async_device_available_fn,
+ device_info_fn=async_device_device_info_fn,
+ name_fn=lambda _: f"{name} {wan} latency",
+ object_fn=lambda api, obj_id: api.devices[obj_id],
+ supported_fn=partial(
+ async_device_wan_latency_supported_fn, wan, monitor_target
+ ),
+ unique_id_fn=lambda hub,
+ obj_id: f"{name.lower}_{wan.lower}_latency-{obj_id}",
+ value_fn=partial(async_device_wan_latency_value_fn, wan, monitor_target),
+ )
+
+ wans: tuple[Literal["WAN"], Literal["WAN2"]] = ("WAN", "WAN2")
+ return tuple(
+ make_wan_latency_entity_description(wan, name, target)
+ for wan in wans
+ for name, target in (
+ ("Microsoft", "microsoft"),
+ ("Google", "google"),
+ ("Cloudflare", "1.1.1.1"),
+ )
+ )
+
+
@dataclass(frozen=True, kw_only=True)
class UnifiSensorEntityDescription(
SensorEntityDescription, UnifiEntityDescription[HandlerT, ApiItemT]
@@ -456,6 +543,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
),
)
+ENTITY_DESCRIPTIONS += make_wan_latency_sensors()
+
async def async_setup_entry(
hass: HomeAssistant,
diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py
index e1893922f60..779df6660f0 100644
--- a/tests/components/unifi/test_sensor.py
+++ b/tests/components/unifi/test_sensor.py
@@ -1424,3 +1424,227 @@ async def test_device_uptime(
entity_registry.async_get("sensor.device_uptime").entity_category
is EntityCategory.DIAGNOSTIC
)
+
+
+@pytest.mark.parametrize(
+ "device_payload",
+ [
+ [
+ {
+ "board_rev": 2,
+ "device_id": "mock-id",
+ "ip": "10.0.1.1",
+ "mac": "10:00:00:00:01:01",
+ "last_seen": 1562600145,
+ "model": "US16P150",
+ "name": "mock-name",
+ "port_overrides": [],
+ "uptime_stats": {
+ "WAN": {
+ "availability": 100.0,
+ "latency_average": 39,
+ "monitors": [
+ {
+ "availability": 100.0,
+ "latency_average": 56,
+ "target": "www.microsoft.com",
+ "type": "icmp",
+ },
+ {
+ "availability": 100.0,
+ "latency_average": 53,
+ "target": "google.com",
+ "type": "icmp",
+ },
+ {
+ "availability": 100.0,
+ "latency_average": 30,
+ "target": "1.1.1.1",
+ "type": "icmp",
+ },
+ ],
+ },
+ "WAN2": {
+ "monitors": [
+ {
+ "availability": 0.0,
+ "target": "www.microsoft.com",
+ "type": "icmp",
+ },
+ {
+ "availability": 0.0,
+ "target": "google.com",
+ "type": "icmp",
+ },
+ {"availability": 0.0, "target": "1.1.1.1", "type": "icmp"},
+ ],
+ },
+ },
+ "state": 1,
+ "type": "usw",
+ "version": "4.0.42.10433",
+ }
+ ]
+ ],
+)
+@pytest.mark.parametrize(
+ ("entity_id", "state", "updated_state", "index_to_update"),
+ [
+ # Microsoft
+ ("microsoft_wan", "56", "20", 0),
+ # Google
+ ("google_wan", "53", "90", 1),
+ # Cloudflare
+ ("cloudflare_wan", "30", "80", 2),
+ ],
+)
+@pytest.mark.usefixtures("config_entry_setup")
+async def test_wan_monitor_latency(
+ hass: HomeAssistant,
+ entity_registry: er.EntityRegistry,
+ mock_websocket_message,
+ device_payload: list[dict[str, Any]],
+ entity_id: str,
+ state: str,
+ updated_state: str,
+ index_to_update: int,
+) -> None:
+ """Verify that wan latency sensors are working as expected."""
+
+ assert len(hass.states.async_all()) == 6
+ assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2
+
+ latency_entry = entity_registry.async_get(f"sensor.mock_name_{entity_id}_latency")
+ assert latency_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
+ assert latency_entry.entity_category is EntityCategory.DIAGNOSTIC
+
+ # Enable entity
+ entity_registry.async_update_entity(
+ entity_id=f"sensor.mock_name_{entity_id}_latency", disabled_by=None
+ )
+
+ await hass.async_block_till_done()
+
+ async_fire_time_changed(
+ hass,
+ dt_util.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1),
+ )
+ await hass.async_block_till_done()
+
+ assert len(hass.states.async_all()) == 7
+ assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3
+
+ # Verify sensor attributes and state
+ latency_entry = hass.states.get(f"sensor.mock_name_{entity_id}_latency")
+ assert latency_entry.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DURATION
+ assert (
+ latency_entry.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT
+ )
+ assert latency_entry.state == state
+
+ # Verify state update
+ device = device_payload[0]
+ device["uptime_stats"]["WAN"]["monitors"][index_to_update]["latency_average"] = (
+ updated_state
+ )
+
+ mock_websocket_message(message=MessageKey.DEVICE, data=device)
+
+ assert (
+ hass.states.get(f"sensor.mock_name_{entity_id}_latency").state == updated_state
+ )
+
+
+@pytest.mark.parametrize(
+ "device_payload",
+ [
+ [
+ {
+ "board_rev": 2,
+ "device_id": "mock-id",
+ "ip": "10.0.1.1",
+ "mac": "10:00:00:00:01:01",
+ "last_seen": 1562600145,
+ "model": "US16P150",
+ "name": "mock-name",
+ "port_overrides": [],
+ "uptime_stats": {
+ "WAN": {
+ "monitors": [
+ {
+ "availability": 100.0,
+ "latency_average": 30,
+ "target": "1.2.3.4",
+ "type": "icmp",
+ },
+ ],
+ },
+ "WAN2": {
+ "monitors": [
+ {
+ "availability": 0.0,
+ "target": "www.microsoft.com",
+ "type": "icmp",
+ },
+ {
+ "availability": 0.0,
+ "target": "google.com",
+ "type": "icmp",
+ },
+ {"availability": 0.0, "target": "1.1.1.1", "type": "icmp"},
+ ],
+ },
+ },
+ "state": 1,
+ "type": "usw",
+ "version": "4.0.42.10433",
+ }
+ ]
+ ],
+)
+@pytest.mark.usefixtures("config_entry_setup")
+async def test_wan_monitor_latency_with_no_entries(
+ hass: HomeAssistant,
+ entity_registry: er.EntityRegistry,
+) -> None:
+ """Verify that wan latency sensors is not created if there is no data."""
+
+ assert len(hass.states.async_all()) == 6
+ assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2
+
+ latency_entry = entity_registry.async_get("sensor.mock_name_google_wan_latency")
+ assert latency_entry is None
+
+
+@pytest.mark.parametrize(
+ "device_payload",
+ [
+ [
+ {
+ "board_rev": 2,
+ "device_id": "mock-id",
+ "ip": "10.0.1.1",
+ "mac": "10:00:00:00:01:01",
+ "last_seen": 1562600145,
+ "model": "US16P150",
+ "name": "mock-name",
+ "port_overrides": [],
+ "state": 1,
+ "type": "usw",
+ "version": "4.0.42.10433",
+ }
+ ]
+ ],
+)
+@pytest.mark.usefixtures("config_entry_setup")
+async def test_wan_monitor_latency_with_no_uptime(
+ hass: HomeAssistant,
+ entity_registry: er.EntityRegistry,
+) -> None:
+ """Verify that wan latency sensors is not created if there is no data."""
+
+ assert len(hass.states.async_all()) == 6
+ assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2
+
+ latency_entry = entity_registry.async_get("sensor.mock_name_google_wan_latency")
+ assert latency_entry is None
--
GitLab