diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 296857e1cfa..fcfe71a2858 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/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, diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index 70b28e34dd0..54b9cb12157 100644 --- a/homeassistant/components/unifi/entity.py +++ b/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. diff --git a/homeassistant/components/unifi/image.py b/homeassistant/components/unifi/image.py index c26f06cb5f2..25c368880fa 100644 --- a/homeassistant/components/unifi/image.py +++ b/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, diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 3682fa0bf6c..8cdc0dcbb71 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/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, + ), ) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index ca11cdfea30..64e3ec2455c 100644 --- a/homeassistant/components/unifi/switch.py +++ b/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}", ), diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index ea02b144a2f..661a9016bdc 100644 --- a/homeassistant/components/unifi/update.py +++ b/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}", diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index bf7ba4d53c0..d619cd4c3c9 100644 --- a/tests/components/unifi/test_sensor.py +++ b/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"