Add UniFi sensor for number of clients connected to a device (#119509)

Co-authored-by: Kim de Vos <[email protected]>

homeassistant/components/unifi/sensor.py

@@ -108,6 +108,27 @@ def async_wlan_client_value_fn(hub: UnifiHub, wlan: Wlan) -> int:
     )
 
 
+@callback
+def async_device_clients_value_fn(hub: UnifiHub, device: Device) -> int:
+    """Calculate the amount of clients connected to a device."""
+
+    return len(
+        [
+            client.mac
+            for client in hub.api.clients.values()
+            if (
+                (
+                    client.access_point_mac != ""
+                    and client.access_point_mac == device.mac
+                )
+                or (client.access_point_mac == "" and client.switch_mac == device.mac)
+            )
+            and dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0)
+            < hub.config.option_detection_time
+        ]
+    )
+
+
 @callback
 def async_device_uptime_value_fn(hub: UnifiHub, device: Device) -> datetime | None:
     """Calculate the approximate time the device started (based on uptime returned from API, in seconds)."""
@@ -302,6 +323,20 @@ class UnifiSensorEntityDescription(
         unique_id_fn=lambda hub, obj_id: f"wlan_clients-{obj_id}",
         value_fn=async_wlan_client_value_fn,
     ),
+    UnifiSensorEntityDescription[Devices, Device](
+        key="Device clients",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        state_class=SensorStateClass.MEASUREMENT,
+        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 device: "Clients",
+        object_fn=lambda api, obj_id: api.devices[obj_id],
+        should_poll=True,
+        unique_id_fn=lambda hub, obj_id: f"device_clients-{obj_id}",
+        value_fn=async_device_clients_value_fn,
+    ),
     UnifiSensorEntityDescription[Outlets, Outlet](
         key="Outlet power metering",
         device_class=SensorDeviceClass.POWER,

tests/components/unifi/test_sensor.py

@@ -1225,3 +1225,110 @@ async def test_bandwidth_port_sensors(
     assert hass.states.get("sensor.mock_name_port_1_tx") is None
     assert hass.states.get("sensor.mock_name_port_2_rx") is None
     assert hass.states.get("sensor.mock_name_port_2_tx") is None
+
+
+@pytest.mark.parametrize(
+    "device_payload",
+    [
+        [
+            {
+                "device_id": "mock-id1",
+                "mac": "01:00:00:00:00:00",
+                "model": "US16P150",
+                "name": "Wired Device",
+                "state": 1,
+                "version": "4.0.42.10433",
+            },
+            {
+                "device_id": "mock-id2",
+                "mac": "02:00:00:00:00:00",
+                "model": "US16P150",
+                "name": "Wireless Device",
+                "state": 1,
+                "version": "4.0.42.10433",
+            },
+        ]
+    ],
+)
+async def test_device_client_sensors(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    config_entry_factory,
+    mock_websocket_message,
+    client_payload,
+) -> None:
+    """Verify that WLAN client sensors are working as expected."""
+    client_payload += [
+        {
+            "hostname": "Wired client 1",
+            "is_wired": True,
+            "mac": "00:00:00:00:00:01",
+            "oui": "Producer",
+            "sw_mac": "01:00:00:00:00:00",
+            "last_seen": dt_util.as_timestamp(dt_util.utcnow()),
+        },
+        {
+            "hostname": "Wired client 2",
+            "is_wired": True,
+            "mac": "00:00:00:00:00:02",
+            "oui": "Producer",
+            "sw_mac": "01:00:00:00:00:00",
+            "last_seen": dt_util.as_timestamp(dt_util.utcnow()),
+        },
+        {
+            "is_wired": False,
+            "mac": "00:00:00:00:00:03",
+            "name": "Wireless client 1",
+            "oui": "Producer",
+            "ap_mac": "02:00:00:00:00:00",
+            "sw_mac": "01:00:00:00:00:00",
+            "last_seen": dt_util.as_timestamp(dt_util.utcnow()),
+        },
+    ]
+    await config_entry_factory()
+
+    assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4
+
+    ent_reg_entry = entity_registry.async_get("sensor.wired_device_clients")
+    assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
+    assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
+    assert ent_reg_entry.unique_id == "device_clients-01:00:00:00:00:00"
+
+    ent_reg_entry = entity_registry.async_get("sensor.wireless_device_clients")
+    assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
+    assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
+    assert ent_reg_entry.unique_id == "device_clients-02:00:00:00:00:00"
+
+    # Enable entity
+    entity_registry.async_update_entity(
+        entity_id="sensor.wired_device_clients", disabled_by=None
+    )
+    entity_registry.async_update_entity(
+        entity_id="sensor.wireless_device_clients", 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()
+
+    # Validate state object
+    assert len(hass.states.async_all()) == 13
+    assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 6
+
+    assert hass.states.get("sensor.wired_device_clients").state == "2"
+    assert hass.states.get("sensor.wireless_device_clients").state == "1"
+
+    # Verify state update - decreasing number
+    wireless_client_1 = client_payload[2]
+    wireless_client_1["last_seen"] = 0
+    mock_websocket_message(message=MessageKey.CLIENT, data=wireless_client_1)
+
+    async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.wired_device_clients").state == "2"
+    assert hass.states.get("sensor.wireless_device_clients").state == "0"