Add PoE power sensor to UniFi integration (#84314)

* Add PoE power sensor to UniFi integration

* Add unit of power

* Only update state if value has changed

* Remove stale print

* Subscribe to specific sensor to remove unnecessary state changes

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/unifi/sensor.py

@@ -8,12 +8,14 @@ from __future__ import annotations
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-from typing import Generic, TypeVar
+from typing import Generic, TypeVar, Union
 
 import aiounifi
 from aiounifi.interfaces.api_handlers import ItemEvent
 from aiounifi.interfaces.clients import Clients
+from aiounifi.interfaces.ports import Ports
 from aiounifi.models.client import Client
+from aiounifi.models.port import Port
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -21,7 +23,7 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import UnitOfInformation
+from homeassistant.const import UnitOfInformation, UnitOfPower
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
@@ -30,11 +32,11 @@ from homeassistant.helpers.entity import DeviceInfo, EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
 
-from .const import DOMAIN as UNIFI_DOMAIN
+from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
 from .controller import UniFiController
 
-_DataT = TypeVar("_DataT", bound=Client)
+_DataT = TypeVar("_DataT", bound=Union[Client, Port])
-_HandlerT = TypeVar("_HandlerT", bound=Clients)
+_HandlerT = TypeVar("_HandlerT", bound=Union[Clients, Ports])
 
 
 @callback
@@ -74,6 +76,31 @@ def async_client_device_info_fn(api: aiounifi.Controller, obj_id: str) -> Device
     )
 
 
+@callback
+def async_device_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
+    """Create device registry entry for device."""
+    if "_" in obj_id:  # Sub device
+        obj_id = obj_id.partition("_")[0]
+
+    device = api.devices[obj_id]
+    return DeviceInfo(
+        connections={(CONNECTION_NETWORK_MAC, device.mac)},
+        manufacturer=ATTR_MANUFACTURER,
+        model=device.model,
+        name=device.name or device.model,
+        sw_version=device.version,
+        hw_version=str(device.board_revision),
+    )
+
+
+@callback
+def async_sub_device_available_fn(controller: UniFiController, obj_id: str) -> bool:
+    """Check if sub device object is disabled."""
+    device_id = obj_id.partition("_")[0]
+    device = controller.api.devices[device_id]
+    return controller.available and not device.disabled
+
+
 @dataclass
 class UnifiEntityLoader(Generic[_HandlerT, _DataT]):
     """Validate and load entities from different UniFi handlers."""
@@ -86,7 +113,7 @@ class UnifiEntityLoader(Generic[_HandlerT, _DataT]):
     object_fn: Callable[[aiounifi.Controller, str], _DataT]
     supported_fn: Callable[[UniFiController, str], bool | None]
     unique_id_fn: Callable[[str], str]
-    value_fn: Callable[[UniFiController, _DataT], datetime | float]
+    value_fn: Callable[[UniFiController, _DataT], datetime | float | str | None]
 
 
 @dataclass
@@ -127,6 +154,23 @@ ENTITY_DESCRIPTIONS: tuple[UnifiEntityDescription, ...] = (
         unique_id_fn=lambda obj_id: f"tx-{obj_id}",
         value_fn=async_client_tx_value_fn,
     ),
+    UnifiEntityDescription[Ports, Port](
+        key="PoE port power sensor",
+        device_class=SensorDeviceClass.POWER,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=UnitOfPower.WATT,
+        has_entity_name=True,
+        entity_registry_enabled_default=False,
+        allowed_fn=lambda controller, obj_id: True,
+        api_handler_fn=lambda api: api.ports,
+        available_fn=async_sub_device_available_fn,
+        device_info_fn=async_device_device_info_fn,
+        name_fn=lambda port: f"{port.name} PoE Power",
+        object_fn=lambda api, obj_id: api.ports[obj_id],
+        supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe,
+        unique_id_fn=lambda obj_id: f"poe_power-{obj_id}",
+        value_fn=lambda _, obj: obj.poe_power if obj.poe_mode != "off" else "0",
+    ),
     UnifiEntityDescription[Clients, Client](
         key="Uptime sensor",
         device_class=SensorDeviceClass.TIMESTAMP,
@@ -253,11 +297,18 @@ class UnifiSensorEntity(SensorEntity, Generic[_HandlerT, _DataT]):
             self.hass.async_create_task(self.remove_item({self._obj_id}))
             return
 
+        update_state = False
         obj = description.object_fn(self.controller.api, self._obj_id)
         if (value := description.value_fn(self.controller, obj)) != self.native_value:
             self._attr_native_value = value
-        self._attr_available = description.available_fn(self.controller, self._obj_id)
+            update_state = True
-        self.async_write_ha_state()
+        if (
+            available := description.available_fn(self.controller, self._obj_id)
+        ) != self.available:
+            self._attr_available = available
+            update_state = True
+        if update_state:
+            self.async_write_ha_state()
 
     @callback
     def async_signal_reachable_callback(self) -> None:

tests/components/unifi/test_sensor.py

@@ -1,25 +1,100 @@
 """UniFi Network sensor platform tests."""
 
-from datetime import datetime
+from copy import deepcopy
+from datetime import datetime, timedelta
 from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
+from aiounifi.websocket import WebsocketState
 import pytest
 
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
-from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN, SensorDeviceClass
 from homeassistant.components.unifi.const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_ALLOW_UPTIME_SENSORS,
     CONF_TRACK_CLIENTS,
     CONF_TRACK_DEVICES,
 )
+from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
+from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity import EntityCategory
+from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 import homeassistant.util.dt as dt_util
 
 from .test_controller import setup_unifi_integration
 
+from tests.common import async_fire_time_changed
+
+DEVICE_1 = {
+    "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": True,
+            "poe_mode": "auto",
+            "poe_power": "2.56",
+            "poe_voltage": "53.40",
+            "portconf_id": "1a1",
+            "port_poe": True,
+            "up": True,
+        },
+        {
+            "media": "GE",
+            "name": "Port 2",
+            "port_idx": 2,
+            "poe_class": "Class 4",
+            "poe_enable": True,
+            "poe_mode": "auto",
+            "poe_power": "2.56",
+            "poe_voltage": "53.40",
+            "portconf_id": "1a2",
+            "port_poe": True,
+            "up": True,
+        },
+        {
+            "media": "GE",
+            "name": "Port 3",
+            "port_idx": 3,
+            "poe_class": "Unknown",
+            "poe_enable": False,
+            "poe_mode": "off",
+            "poe_power": "0.00",
+            "poe_voltage": "0.00",
+            "portconf_id": "1a3",
+            "port_poe": False,
+            "up": True,
+        },
+        {
+            "media": "GE",
+            "name": "Port 4",
+            "port_idx": 4,
+            "poe_class": "Unknown",
+            "poe_enable": False,
+            "poe_mode": "auto",
+            "poe_power": "0.00",
+            "poe_voltage": "0.00",
+            "portconf_id": "1a4",
+            "port_poe": True,
+            "up": True,
+        },
+    ],
+    "state": 1,
+    "type": "usw",
+    "version": "4.0.42.10433",
+}
+
 
 async def test_no_clients(hass, aioclient_mock):
     """Test the update_clients function when no clients are found."""
@@ -243,3 +318,73 @@ async def test_remove_sensors(hass, aioclient_mock, mock_unifi_websocket):
     assert hass.states.get("sensor.wireless_client_rx")
     assert hass.states.get("sensor.wireless_client_tx")
     assert hass.states.get("sensor.wireless_client_uptime")
+
+
+async def test_poe_port_switches(hass, aioclient_mock, mock_unifi_websocket):
+    """Test the update_items function with some clients."""
+    await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1])
+    assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 0
+
+    ent_reg = er.async_get(hass)
+    ent_reg_entry = ent_reg.async_get("sensor.mock_name_port_1_poe_power")
+    assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
+    assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
+
+    # Enable entity
+    ent_reg.async_update_entity(
+        entity_id="sensor.mock_name_port_1_poe_power", 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
+    poe_sensor = hass.states.get("sensor.mock_name_port_1_poe_power")
+    assert poe_sensor.state == "2.56"
+    assert poe_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.POWER
+
+    # Update state object
+    device_1 = deepcopy(DEVICE_1)
+    device_1["port_table"][0]["poe_power"] = "5.12"
+    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").state == "5.12"
+
+    # PoE is disabled
+    device_1 = deepcopy(DEVICE_1)
+    device_1["port_table"][0]["poe_mode"] = "off"
+    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").state == "0"
+
+    # Availability signalling
+
+    # Controller disconnects
+    mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
+    await hass.async_block_till_done()
+    assert (
+        hass.states.get("sensor.mock_name_port_1_poe_power").state == STATE_UNAVAILABLE
+    )
+
+    # Controller reconnects
+    mock_unifi_websocket(state=WebsocketState.RUNNING)
+    await hass.async_block_till_done()
+    assert hass.states.get("sensor.mock_name_port_1_poe_power")
+
+    # Device gets disabled
+    device_1["disabled"] = True
+    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").state == STATE_UNAVAILABLE
+    )
+
+    # Device gets re-enabled
+    device_1["disabled"] = False
+    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")