From b69f589c30e3e1908f26edd3c9103a13ab1d8417 Mon Sep 17 00:00:00 2001 From: Kim de Vos Date: Mon, 22 Apr 2024 22:39:46 +0200 Subject: [PATCH] Add bandwidth sensor for unifi device ports (#115362) --- homeassistant/components/unifi/sensor.py | 37 ++++++ tests/components/unifi/test_sensor.py | 149 +++++++++++++++++++++++ 2 files changed, 186 insertions(+) diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 7d9720cde1a..3979f45ecd8 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/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 cast from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.clients import Clients @@ -239,6 +240,42 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = ( unique_id_fn=lambda hub, obj_id: f"poe_power-{obj_id}", value_fn=lambda _, obj: obj.poe_power if obj.poe_mode != "off" else "0", ), + UnifiSensorEntityDescription[Ports, Port]( + key="Port Bandwidth sensor RX", + device_class=SensorDeviceClass.DATA_RATE, + entity_category=EntityCategory.DIAGNOSTIC, + entity_registry_enabled_default=False, + state_class=SensorStateClass.MEASUREMENT, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + suggested_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND, + icon="mdi:download", + allowed_fn=lambda hub, _: hub.config.option_allow_bandwidth_sensors, + api_handler_fn=lambda api: api.ports, + available_fn=async_device_available_fn, + device_info_fn=async_device_device_info_fn, + name_fn=lambda port: f"{port.name} RX", + object_fn=lambda api, obj_id: api.ports[obj_id], + unique_id_fn=lambda hub, obj_id: f"port_rx-{obj_id}", + value_fn=lambda hub, port: cast(float, port.raw.get("rx_bytes-r", 0)), + ), + UnifiSensorEntityDescription[Ports, Port]( + key="Port Bandwidth sensor TX", + device_class=SensorDeviceClass.DATA_RATE, + entity_category=EntityCategory.DIAGNOSTIC, + entity_registry_enabled_default=False, + state_class=SensorStateClass.MEASUREMENT, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + suggested_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND, + icon="mdi:upload", + allowed_fn=lambda hub, _: hub.config.option_allow_bandwidth_sensors, + api_handler_fn=lambda api: api.ports, + available_fn=async_device_available_fn, + device_info_fn=async_device_device_info_fn, + name_fn=lambda port: f"{port.name} TX", + object_fn=lambda api, obj_id: api.ports[obj_id], + unique_id_fn=lambda hub, obj_id: f"port_tx-{obj_id}", + value_fn=lambda hub, port: cast(float, port.raw.get("tx_bytes-r", 0)), + ), UnifiSensorEntityDescription[Clients, Client]( key="Client uptime", device_class=SensorDeviceClass.TIMESTAMP, diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index e3b4ddd3b63..26eadfa498e 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -1042,3 +1042,152 @@ async def test_device_system_stats( assert hass.states.get("sensor.device_cpu_utilization").state == "7.7" assert hass.states.get("sensor.device_memory_utilization").state == "33.3" + + +async def test_bandwidth_port_sensors( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + mock_unifi_websocket, +) -> None: + """Verify that port bandwidth sensors are working as expected.""" + device_reponse = { + "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": [], + "port_table": [ + { + "media": "GE", + "name": "Port 1", + "port_idx": 1, + "poe_class": "Class 4", + "poe_enable": False, + "poe_mode": "auto", + "poe_power": "2.56", + "poe_voltage": "53.40", + "portconf_id": "1a1", + "port_poe": False, + "up": True, + "rx_bytes-r": 1151, + "tx_bytes-r": 5111, + }, + { + "media": "GE", + "name": "Port 2", + "port_idx": 2, + "poe_class": "Class 4", + "poe_enable": False, + "poe_mode": "auto", + "poe_power": "2.56", + "poe_voltage": "53.40", + "portconf_id": "1a2", + "port_poe": False, + "up": True, + "rx_bytes-r": 1536, + "tx_bytes-r": 3615, + }, + ], + "state": 1, + "type": "usw", + "version": "4.0.42.10433", + } + options = { + CONF_ALLOW_BANDWIDTH_SENSORS: True, + CONF_ALLOW_UPTIME_SENSORS: False, + CONF_TRACK_CLIENTS: False, + CONF_TRACK_DEVICES: False, + } + + config_entry = await setup_unifi_integration( + hass, + aioclient_mock, + options=options, + devices_response=[device_reponse], + ) + + assert len(hass.states.async_all()) == 5 + assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2 + + p1rx_reg_entry = entity_registry.async_get("sensor.mock_name_port_1_rx") + assert p1rx_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION + assert p1rx_reg_entry.entity_category is EntityCategory.DIAGNOSTIC + + p1tx_reg_entry = entity_registry.async_get("sensor.mock_name_port_1_tx") + assert p1tx_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION + assert p1tx_reg_entry.entity_category is EntityCategory.DIAGNOSTIC + + # Enable entity + entity_registry.async_update_entity( + entity_id="sensor.mock_name_port_1_rx", disabled_by=None + ) + entity_registry.async_update_entity( + entity_id="sensor.mock_name_port_1_tx", disabled_by=None + ) + entity_registry.async_update_entity( + entity_id="sensor.mock_name_port_2_rx", disabled_by=None + ) + entity_registry.async_update_entity( + entity_id="sensor.mock_name_port_2_tx", 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()) == 9 + assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 6 + + # Verify sensor attributes and state + p1rx_sensor = hass.states.get("sensor.mock_name_port_1_rx") + assert p1rx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert p1rx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert p1rx_sensor.state == "0.00921" + + p1tx_sensor = hass.states.get("sensor.mock_name_port_1_tx") + assert p1tx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert p1tx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert p1tx_sensor.state == "0.04089" + + p2rx_sensor = hass.states.get("sensor.mock_name_port_2_rx") + assert p2rx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert p2rx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert p2rx_sensor.state == "0.01229" + + p2tx_sensor = hass.states.get("sensor.mock_name_port_2_tx") + assert p2tx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert p2tx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert p2tx_sensor.state == "0.02892" + + # Verify state update + device_reponse["port_table"][0]["rx_bytes-r"] = 3456000000 + device_reponse["port_table"][0]["tx_bytes-r"] = 7891000000 + + mock_unifi_websocket(message=MessageKey.DEVICE, data=device_reponse) + await hass.async_block_till_done() + + assert hass.states.get("sensor.mock_name_port_1_rx").state == "27648.00000" + assert hass.states.get("sensor.mock_name_port_1_tx").state == "63128.00000" + + # Disable option + options[CONF_ALLOW_BANDWIDTH_SENSORS] = False + hass.config_entries.async_update_entry(config_entry, options=options.copy()) + await hass.async_block_till_done() + + assert len(hass.states.async_all()) == 5 + assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2 + + assert hass.states.get("sensor.mock_name_uptime") + assert hass.states.get("sensor.mock_name_state") + assert hass.states.get("sensor.mock_name_port_1_rx") is None + 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