Add bandwidth sensor for unifi device ports (#115362)

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,

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