From 707e422a3167eec2f7a5d4aa765c181c80874085 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Wed, 12 Jun 2024 18:20:31 +0200 Subject: [PATCH] Add UniFi sensor for number of clients connected to a device (#119509) Co-authored-by: Kim de Vos --- homeassistant/components/unifi/sensor.py | 35 ++++++++ tests/components/unifi/test_sensor.py | 107 +++++++++++++++++++++++ 2 files changed, 142 insertions(+) diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 3fd179f5676..ba1da7ea6c8 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/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 @@ ENTITY_DESCRIPTIONS: tuple[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, diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 735df53b0c5..802166068b2 100644 --- a/tests/components/unifi/test_sensor.py +++ b/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"