Add WLAN clients reporting to UniFi Sensor platform (#97234)

homeassistant/components/unifi/device_tracker.py

@@ -171,6 +171,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = (
         is_connected_fn=async_client_is_connected_fn,
         name_fn=lambda client: client.name or client.hostname,
         object_fn=lambda api, obj_id: api.clients[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, obj_id: True,
         unique_id_fn=lambda controller, obj_id: f"{obj_id}-{controller.site}",
         ip_address_fn=lambda api, obj_id: api.clients[obj_id].ip,
@@ -190,6 +191,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = (
         is_connected_fn=lambda ctrlr, obj_id: ctrlr.api.devices[obj_id].state == 1,
         name_fn=lambda device: device.name or device.model,
         object_fn=lambda api, obj_id: api.devices[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, obj_id: True,
         unique_id_fn=lambda controller, obj_id: obj_id,
         ip_address_fn=lambda api, obj_id: api.devices[obj_id].ip,

homeassistant/components/unifi/entity.py

@@ -86,6 +86,7 @@ class UnifiDescription(Generic[HandlerT, ApiItemT]):
     event_to_subscribe: tuple[EventKey, ...] | None
     name_fn: Callable[[ApiItemT], str | None]
     object_fn: Callable[[aiounifi.Controller, str], ApiItemT]
+    should_poll: bool
     supported_fn: Callable[[UniFiController, str], bool | None]
     unique_id_fn: Callable[[UniFiController, str], str]
 
@@ -99,8 +100,6 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
     """Representation of a UniFi entity."""
 
     entity_description: UnifiEntityDescription[HandlerT, ApiItemT]
-    _attr_should_poll = False
-
     _attr_unique_id: str
 
     def __init__(
@@ -120,6 +119,7 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
 
         self._attr_available = description.available_fn(controller, obj_id)
         self._attr_device_info = description.device_info_fn(controller.api, obj_id)
+        self._attr_should_poll = description.should_poll
         self._attr_unique_id = description.unique_id_fn(controller, obj_id)
 
         obj = description.object_fn(self.controller.api, obj_id)
@@ -209,6 +209,10 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
         else:
             await self.async_remove(force_remove=True)
 
+    async def async_update(self) -> None:
+        """Update state if polling is configured."""
+        self.async_update_state(ItemEvent.CHANGED, self._obj_id)
+
     @callback
     def async_initiate_state(self) -> None:
         """Initiate entity state.

homeassistant/components/unifi/image.py

@@ -67,6 +67,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiImageEntityDescription, ...] = (
         event_to_subscribe=None,
         name_fn=lambda _: "QR Code",
         object_fn=lambda api, obj_id: api.wlans[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, obj_id: True,
         unique_id_fn=lambda controller, obj_id: f"qr_code-{obj_id}",
         image_fn=async_wlan_qr_code_image_fn,

homeassistant/components/unifi/sensor.py

@@ -14,9 +14,11 @@ import aiounifi
 from aiounifi.interfaces.api_handlers import ItemEvent
 from aiounifi.interfaces.clients import Clients
 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.port import Port
+from aiounifi.models.wlan import Wlan
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -39,6 +41,7 @@ from .entity import (
     UnifiEntityDescription,
     async_device_available_fn,
     async_device_device_info_fn,
+    async_wlan_device_info_fn,
 )
 
 
@@ -68,6 +71,18 @@ def async_client_uptime_value_fn(
     return dt_util.utc_from_timestamp(float(client.uptime))
 
 
+@callback
+def async_wlan_client_value_fn(controller: UniFiController, wlan: Wlan) -> int:
+    """Calculate the amount of clients connected to a wlan."""
+    return len(
+        [
+            client.mac
+            for client in controller.api.clients.values()
+            if client.essid == wlan.name
+        ]
+    )
+
+
 @callback
 def async_client_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
     """Create device registry entry for client."""
@@ -109,6 +124,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         event_to_subscribe=None,
         name_fn=lambda _: "RX",
         object_fn=lambda api, obj_id: api.clients[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, _: controller.option_allow_bandwidth_sensors,
         unique_id_fn=lambda controller, obj_id: f"rx-{obj_id}",
         value_fn=async_client_rx_value_fn,
@@ -126,6 +142,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         event_to_subscribe=None,
         name_fn=lambda _: "TX",
         object_fn=lambda api, obj_id: api.clients[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, _: controller.option_allow_bandwidth_sensors,
         unique_id_fn=lambda controller, obj_id: f"tx-{obj_id}",
         value_fn=async_client_tx_value_fn,
@@ -145,6 +162,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         event_to_subscribe=None,
         name_fn=lambda port: f"{port.name} PoE Power",
         object_fn=lambda api, obj_id: api.ports[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe,
         unique_id_fn=lambda controller, obj_id: f"poe_power-{obj_id}",
         value_fn=lambda _, obj: obj.poe_power if obj.poe_mode != "off" else "0",
@@ -163,10 +181,28 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         event_to_subscribe=None,
         name_fn=lambda client: "Uptime",
         object_fn=lambda api, obj_id: api.clients[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, _: controller.option_allow_uptime_sensors,
         unique_id_fn=lambda controller, obj_id: f"uptime-{obj_id}",
         value_fn=async_client_uptime_value_fn,
     ),
+    UnifiSensorEntityDescription[Wlans, Wlan](
+        key="WLAN clients",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        has_entity_name=True,
+        allowed_fn=lambda controller, _: True,
+        api_handler_fn=lambda api: api.wlans,
+        available_fn=lambda controller, obj_id: controller.available,
+        device_info_fn=async_wlan_device_info_fn,
+        event_is_on=None,
+        event_to_subscribe=None,
+        name_fn=lambda client: None,
+        object_fn=lambda api, obj_id: api.wlans[obj_id],
+        should_poll=True,
+        supported_fn=lambda controller, _: True,
+        unique_id_fn=lambda controller, obj_id: f"wlan_clients-{obj_id}",
+        value_fn=async_wlan_client_value_fn,
+    ),
 )
 
 

homeassistant/components/unifi/switch.py

@@ -186,6 +186,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
         name_fn=lambda client: None,
         object_fn=lambda api, obj_id: api.clients[obj_id],
         only_event_for_state_change=True,
+        should_poll=False,
         supported_fn=lambda controller, obj_id: True,
         unique_id_fn=lambda controller, obj_id: f"block-{obj_id}",
     ),
@@ -204,6 +205,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
         is_on_fn=async_dpi_group_is_on_fn,
         name_fn=lambda group: group.name,
         object_fn=lambda api, obj_id: api.dpi_groups[obj_id],
+        should_poll=False,
         supported_fn=lambda c, obj_id: bool(c.api.dpi_groups[obj_id].dpiapp_ids),
         unique_id_fn=lambda controller, obj_id: obj_id,
     ),
@@ -221,6 +223,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
         is_on_fn=lambda controller, outlet: outlet.relay_state,
         name_fn=lambda outlet: outlet.name,
         object_fn=lambda api, obj_id: api.outlets[obj_id],
+        should_poll=False,
         supported_fn=lambda c, obj_id: c.api.outlets[obj_id].has_relay,
         unique_id_fn=lambda controller, obj_id: f"{obj_id.split('_', 1)[0]}-outlet-{obj_id.split('_', 1)[1]}",
     ),
@@ -241,6 +244,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
         is_on_fn=lambda controller, port: port.poe_mode != "off",
         name_fn=lambda port: f"{port.name} PoE",
         object_fn=lambda api, obj_id: api.ports[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe,
         unique_id_fn=lambda controller, obj_id: f"{obj_id.split('_', 1)[0]}-poe-{obj_id.split('_', 1)[1]}",
     ),
@@ -260,6 +264,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
         is_on_fn=lambda controller, wlan: wlan.enabled,
         name_fn=lambda wlan: None,
         object_fn=lambda api, obj_id: api.wlans[obj_id],
+        should_poll=False,
         supported_fn=lambda controller, obj_id: True,
         unique_id_fn=lambda controller, obj_id: f"wlan-{obj_id}",
     ),

homeassistant/components/unifi/update.py

@@ -74,6 +74,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiUpdateEntityDescription, ...] = (
         event_to_subscribe=None,
         name_fn=lambda device: None,
         object_fn=lambda api, obj_id: api.devices[obj_id],
+        should_poll=False,
         state_fn=lambda api, device: device.state == 4,
         supported_fn=lambda controller, obj_id: True,
         unique_id_fn=lambda controller, obj_id: f"device_update-{obj_id}",

tests/components/unifi/test_sensor.py

@@ -19,6 +19,7 @@ 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.entity_component import DEFAULT_SCAN_INTERVAL
 from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 import homeassistant.util.dt as dt_util
 
@@ -95,6 +96,42 @@ DEVICE_1 = {
     "version": "4.0.42.10433",
 }
 
+WLAN = {
+    "_id": "012345678910111213141516",
+    "bc_filter_enabled": False,
+    "bc_filter_list": [],
+    "dtim_mode": "default",
+    "dtim_na": 1,
+    "dtim_ng": 1,
+    "enabled": True,
+    "group_rekey": 3600,
+    "mac_filter_enabled": False,
+    "mac_filter_list": [],
+    "mac_filter_policy": "allow",
+    "minrate_na_advertising_rates": False,
+    "minrate_na_beacon_rate_kbps": 6000,
+    "minrate_na_data_rate_kbps": 6000,
+    "minrate_na_enabled": False,
+    "minrate_na_mgmt_rate_kbps": 6000,
+    "minrate_ng_advertising_rates": False,
+    "minrate_ng_beacon_rate_kbps": 1000,
+    "minrate_ng_data_rate_kbps": 1000,
+    "minrate_ng_enabled": False,
+    "minrate_ng_mgmt_rate_kbps": 1000,
+    "name": "SSID 1",
+    "no2ghz_oui": False,
+    "schedule": [],
+    "security": "wpapsk",
+    "site_id": "5a32aa4ee4b0412345678910",
+    "usergroup_id": "012345678910111213141518",
+    "wep_idx": 1,
+    "wlangroup_id": "012345678910111213141519",
+    "wpa_enc": "ccmp",
+    "wpa_mode": "wpa2",
+    "x_iapp_key": "01234567891011121314151617181920",
+    "x_passphrase": "password",
+}
+
 
 async def test_no_clients(
     hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
@@ -424,3 +461,90 @@ async def test_poe_port_switches(
     mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
     await hass.async_block_till_done()
     assert hass.states.get("sensor.mock_name_port_1_poe_power")
+
+
+async def test_wlan_client_sensors(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+) -> None:
+    """Verify that WLAN client sensors are working as expected."""
+    wireless_client_1 = {
+        "essid": "SSID 1",
+        "is_wired": False,
+        "mac": "00:00:00:00:00:01",
+        "name": "Wireless client",
+        "oui": "Producer",
+        "rx_bytes-r": 2345000000,
+        "tx_bytes-r": 6789000000,
+    }
+    wireless_client_2 = {
+        "essid": "SSID 2",
+        "is_wired": False,
+        "mac": "00:00:00:00:00:02",
+        "name": "Wireless client2",
+        "oui": "Producer2",
+        "rx_bytes-r": 2345000000,
+        "tx_bytes-r": 6789000000,
+    }
+
+    await setup_unifi_integration(
+        hass,
+        aioclient_mock,
+        clients_response=[wireless_client_1, wireless_client_2],
+        wlans_response=[WLAN],
+    )
+
+    assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
+
+    ent_reg = er.async_get(hass)
+    ent_reg_entry = ent_reg.async_get("sensor.ssid_1")
+    assert ent_reg_entry.unique_id == "wlan_clients-012345678910111213141516"
+    assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
+
+    # Validate state object
+    ssid_1 = hass.states.get("sensor.ssid_1")
+    assert ssid_1 is not None
+    assert ssid_1.state == "1"
+
+    # Verify state update - increasing number
+
+    wireless_client_1["essid"] = "SSID 1"
+    wireless_client_2["essid"] = "SSID 1"
+
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=wireless_client_1)
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=wireless_client_2)
+    await hass.async_block_till_done()
+
+    ssid_1 = hass.states.get("sensor.ssid_1")
+    assert ssid_1.state == "1"
+
+    async_fire_time_changed(hass, datetime.utcnow() + DEFAULT_SCAN_INTERVAL)
+    await hass.async_block_till_done()
+
+    ssid_1 = hass.states.get("sensor.ssid_1")
+    assert ssid_1.state == "2"
+
+    # Verify state update - decreasing number
+
+    wireless_client_1["essid"] = "SSID"
+    wireless_client_2["essid"] = "SSID"
+
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=wireless_client_1)
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=wireless_client_2)
+
+    async_fire_time_changed(hass, datetime.utcnow() + DEFAULT_SCAN_INTERVAL)
+    await hass.async_block_till_done()
+
+    ssid_1 = hass.states.get("sensor.ssid_1")
+    assert ssid_1.state == "0"
+
+    # Availability signalling
+
+    # Controller disconnects
+    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await hass.async_block_till_done()
+    assert hass.states.get("sensor.ssid_1").state == STATE_UNAVAILABLE
+
+    # Controller reconnects
+    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await hass.async_block_till_done()
+    assert hass.states.get("sensor.ssid_1").state == "0"