From fe1f617a354fdb41453363d8ae8f97e51d08f39d Mon Sep 17 00:00:00 2001
From: Chris <Cisien@users.noreply.github.com>
Date: Thu, 10 Aug 2023 09:25:03 -0700
Subject: [PATCH] Add unifi power stats for PDU outlets (#98081)

Adds support for power stats for PDU outlets.
---
 homeassistant/components/unifi/sensor.py |  31 ++++
 tests/components/unifi/test_sensor.py    | 176 +++++++++++++++++++++++
 2 files changed, 207 insertions(+)

diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py
index 23cc8724c2c..367ff1332f4 100644
--- a/homeassistant/components/unifi/sensor.py
+++ b/homeassistant/components/unifi/sensor.py
@@ -12,10 +12,12 @@ from typing import Generic
 
 from aiounifi.interfaces.api_handlers import ItemEvent
 from aiounifi.interfaces.clients import Clients
+from aiounifi.interfaces.outlets import Outlets
 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.outlet import Outlet
 from aiounifi.models.port import Port
 from aiounifi.models.wlan import Wlan
 
@@ -84,6 +86,16 @@ def async_wlan_client_value_fn(controller: UniFiController, wlan: Wlan) -> int:
     )
 
 
+@callback
+def async_device_outlet_power_supported_fn(
+    controller: UniFiController, obj_id: str
+) -> bool:
+    """Determine if an outlet has the power property."""
+    # At this time, an outlet_caps value of 3 is expected to indicate that the outlet
+    # supports metering
+    return controller.api.outlets[obj_id].caps == 3
+
+
 @dataclass
 class UnifiSensorEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
     """Validate and load entities from different UniFi handlers."""
@@ -193,6 +205,25 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         unique_id_fn=lambda controller, obj_id: f"wlan_clients-{obj_id}",
         value_fn=async_wlan_client_value_fn,
     ),
+    UnifiSensorEntityDescription[Outlets, Outlet](
+        key="Outlet power metering",
+        device_class=SensorDeviceClass.POWER,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=UnitOfPower.WATT,
+        has_entity_name=True,
+        allowed_fn=lambda controller, obj_id: True,
+        api_handler_fn=lambda api: api.outlets,
+        available_fn=async_device_available_fn,
+        device_info_fn=async_device_device_info_fn,
+        event_is_on=None,
+        event_to_subscribe=None,
+        name_fn=lambda outlet: f"{outlet.name} Outlet Power",
+        object_fn=lambda api, obj_id: api.outlets[obj_id],
+        should_poll=True,
+        supported_fn=async_device_outlet_power_supported_fn,
+        unique_id_fn=lambda controller, obj_id: f"outlet_power-{obj_id}",
+        value_fn=lambda _, obj: obj.power if obj.relay_state else "0",
+    ),
 )
 
 
diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py
index 9670ecb43d0..98a4941caaa 100644
--- a/tests/components/unifi/test_sensor.py
+++ b/tests/components/unifi/test_sensor.py
@@ -132,6 +132,152 @@ WLAN = {
     "x_passphrase": "password",
 }
 
+PDU_DEVICE_1 = {
+    "_id": "123456654321abcdef012345",
+    "required_version": "5.28.0",
+    "port_table": [],
+    "license_state": "registered",
+    "lcm_brightness_override": False,
+    "type": "usw",
+    "board_rev": 4,
+    "hw_caps": 136,
+    "reboot_duration": 70,
+    "snmp_contact": "",
+    "config_network": {"type": "dhcp", "bonding_enabled": False},
+    "outlet_table": [
+        {
+            "index": 1,
+            "relay_state": True,
+            "cycle_enabled": False,
+            "name": "USB Outlet 1",
+            "outlet_caps": 1,
+        },
+        {
+            "index": 2,
+            "relay_state": True,
+            "cycle_enabled": False,
+            "name": "Outlet 2",
+            "outlet_caps": 3,
+            "outlet_voltage": "119.644",
+            "outlet_current": "0.935",
+            "outlet_power": "73.827",
+            "outlet_power_factor": "0.659",
+        },
+    ],
+    "model": "USPPDUP",
+    "manufacturer_id": 4,
+    "ip": "192.168.1.76",
+    "fw2_caps": 0,
+    "jumboframe_enabled": False,
+    "version": "6.5.59.14777",
+    "unsupported_reason": 0,
+    "adoption_completed": True,
+    "outlet_enabled": True,
+    "stp_version": "rstp",
+    "name": "Dummy USP-PDU-Pro",
+    "fw_caps": 1732968229,
+    "lcm_brightness": 80,
+    "internet": True,
+    "mgmt_network_id": "123456654321abcdef012347",
+    "gateway_mac": "01:02:03:04:05:06",
+    "stp_priority": "32768",
+    "lcm_night_mode_begins": "22:00",
+    "two_phase_adopt": False,
+    "connected_at": 1690626493,
+    "inform_ip": "192.168.1.1",
+    "cfgversion": "ba8f30a5a17aad64",
+    "mac": "01:02:03:04:05:ff",
+    "provisioned_at": 1690989511,
+    "inform_url": "http://192.168.1.1:8080/inform",
+    "upgrade_duration": 100,
+    "ethernet_table": [{"num_port": 1, "name": "eth0", "mac": "01:02:03:04:05:a1"}],
+    "flowctrl_enabled": False,
+    "unsupported": False,
+    "ble_caps": 0,
+    "sys_error_caps": 0,
+    "dot1x_portctrl_enabled": False,
+    "last_uplink": {},
+    "disconnected_at": 1690626452,
+    "architecture": "mips",
+    "x_aes_gcm": True,
+    "has_fan": False,
+    "outlet_overrides": [
+        {
+            "cycle_enabled": False,
+            "name": "USB Outlet 1",
+            "relay_state": True,
+            "index": 1,
+        },
+        {"cycle_enabled": False, "name": "Outlet 2", "relay_state": True, "index": 2},
+    ],
+    "model_incompatible": False,
+    "satisfaction": 100,
+    "model_in_eol": False,
+    "anomalies": -1,
+    "has_temperature": False,
+    "switch_caps": {},
+    "adopted_by_client": "web",
+    "snmp_location": "",
+    "model_in_lts": False,
+    "kernel_version": "4.14.115",
+    "serial": "abc123",
+    "power_source_ctrl_enabled": False,
+    "lcm_night_mode_ends": "08:00",
+    "adopted": True,
+    "hash_id": "abcdef123456",
+    "device_id": "mock-pdu",
+    "uplink": {},
+    "state": 1,
+    "start_disconnected_millis": 1690626383386,
+    "credential_caps": 0,
+    "default": False,
+    "discovered_via": "l2",
+    "adopt_ip": "10.0.10.4",
+    "adopt_url": "http://192.168.1.1:8080/inform",
+    "last_seen": 1691518814,
+    "min_inform_interval_seconds": 10,
+    "upgradable": False,
+    "adoptable_when_upgraded": False,
+    "rollupgrade": False,
+    "known_cfgversion": "abcfde03929",
+    "uptime": 1193042,
+    "_uptime": 1193042,
+    "locating": False,
+    "start_connected_millis": 1690626493324,
+    "prev_non_busy_state": 5,
+    "next_interval": 47,
+    "sys_stats": {},
+    "system-stats": {"cpu": "1.4", "mem": "28.9", "uptime": "1193042"},
+    "ssh_session_table": [],
+    "lldp_table": [],
+    "displayable_version": "6.5.59",
+    "connection_network_id": "123456654321abcdef012349",
+    "connection_network_name": "Default",
+    "startup_timestamp": 1690325774,
+    "is_access_point": False,
+    "safe_for_autoupgrade": True,
+    "overheating": False,
+    "power_source": "0",
+    "total_max_power": 0,
+    "outlet_ac_power_budget": "1875.000",
+    "outlet_ac_power_consumption": "201.683",
+    "downlink_table": [],
+    "uplink_depth": 1,
+    "downlink_lldp_macs": [],
+    "dhcp_server_table": [],
+    "connect_request_ip": "10.0.10.4",
+    "connect_request_port": "57951",
+    "ipv4_lease_expiration_timestamp_seconds": 1691576686,
+    "stat": {},
+    "tx_bytes": 1426780,
+    "rx_bytes": 1435064,
+    "bytes": 2861844,
+    "num_sta": 0,
+    "user-num_sta": 0,
+    "guest-num_sta": 0,
+    "x_has_ssh_hostkey": True,
+}
+
 
 async def test_no_clients(
     hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
@@ -571,3 +717,33 @@ async def test_wlan_client_sensors(
     mock_unifi_websocket(message=MessageKey.WLAN_CONF_UPDATED, data=wlan_1)
     await hass.async_block_till_done()
     assert hass.states.get("sensor.ssid_1").state == "0"
+
+
+async def test_outlet_power_readings(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
+) -> None:
+    """Test the outlet power reporting on PDU devices."""
+    await setup_unifi_integration(hass, aioclient_mock, devices_response=[PDU_DEVICE_1])
+
+    assert len(hass.states.async_all()) == 5
+    assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
+
+    ent_reg = er.async_get(hass)
+    ent_reg_entry = ent_reg.async_get("sensor.dummy_usp_pdu_pro_outlet_2_outlet_power")
+    assert ent_reg_entry.unique_id == "outlet_power-01:02:03:04:05:ff_2"
+    assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
+
+    outlet_2 = hass.states.get("sensor.dummy_usp_pdu_pro_outlet_2_outlet_power")
+    assert outlet_2.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.POWER
+    assert outlet_2.state == "73.827"
+
+    # Verify state update
+    pdu_device_state_update = deepcopy(PDU_DEVICE_1)
+
+    pdu_device_state_update["outlet_table"][1]["outlet_power"] = "123.45"
+
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=pdu_device_state_update)
+    await hass.async_block_till_done()
+
+    outlet_2 = hass.states.get("sensor.dummy_usp_pdu_pro_outlet_2_outlet_power")
+    assert outlet_2.state == "123.45"
-- 
GitLab