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

@@ -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,

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