Add power monitor to WiZ

Requires sbidy/pywizlight#147

homeassistant/components/wiz/__init__.py

@@ -1,4 +1,6 @@
 """WiZ Platform integration."""
+from __future__ import annotations
+
 import asyncio
 from datetime import timedelta
 import logging
@@ -80,10 +82,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             "Found bulb {bulb.mac} at {ip_address}, expected {entry.unique_id}"
         )
 
-    async def _async_update() -> None:
+    async def _async_update() -> float | None:
         """Update the WiZ device."""
         try:
             await bulb.updateState()
+            if bulb.power_monitoring is not False:
+                power: float | None = await bulb.get_power()
+                return power
+            return None
         except WIZ_EXCEPTIONS as ex:
             raise UpdateFailed(f"Failed to update device at {ip_address}: {ex}") from ex
 
@@ -117,7 +123,7 @@ async def _async_shutdown_on_stop(event: Event) -> None:
     def _async_push_update(state: PilotParser) -> None:
         """Receive a push update."""
         _LOGGER.debug("%s: Got push update: %s", bulb.mac, state.pilotResult)
-        coordinator.async_set_updated_data(None)
+        coordinator.async_set_updated_data(coordinator.data)
         if state.get_source() == PIR_SOURCE:
             async_dispatcher_send(hass, SIGNAL_WIZ_PIR.format(bulb.mac))
 

homeassistant/components/wiz/entity.py

@@ -2,20 +2,23 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import Any
+from typing import Any, Optional
 
 from pywizlight.bulblibrary import BulbType
 
 from homeassistant.const import ATTR_HW_VERSION, ATTR_MODEL
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.entity import DeviceInfo, Entity, ToggleEntity
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
 
 from .models import WizData
 
 
-class WizEntity(CoordinatorEntity, Entity):
+class WizEntity(CoordinatorEntity[DataUpdateCoordinator[Optional[float]]], Entity):
     """Representation of WiZ entity."""
 
     def __init__(self, wiz_data: WizData, name: str) -> None:

homeassistant/components/wiz/sensor.py

@@ -8,7 +8,7 @@
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import SIGNAL_STRENGTH_DECIBELS_MILLIWATT
+from homeassistant.const import POWER_WATT, SIGNAL_STRENGTH_DECIBELS_MILLIWATT
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -30,16 +30,35 @@
 )
 
 
+POWER_SENSORS: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="power",
+        name="Current Power",
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.POWER,
+        native_unit_of_measurement=POWER_WATT,
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: ConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the wiz sensor."""
     wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities(
+    entities = [
         WizSensor(wiz_data, entry.title, description) for description in SENSORS
-    )
+    ]
+    if wiz_data.coordinator.data is not None:
+        entities.extend(
+            [
+                WizPowerSensor(wiz_data, entry.title, description)
+                for description in POWER_SENSORS
+            ]
+        )
+    async_add_entities(entities)
 
 
 class WizSensor(WizEntity, SensorEntity):
@@ -63,3 +82,12 @@ def _async_update_attrs(self) -> None:
         self._attr_native_value = self._device.state.pilotResult.get(
             self.entity_description.key
         )
+
+
+class WizPowerSensor(WizSensor):
+    """Defines a WiZ power sensor."""
+
+    @callback
+    def _async_update_attrs(self) -> None:
+        """Handle updating _attr values."""
+        self._attr_native_value = self.coordinator.data

tests/components/wiz/__init__.py

@@ -150,6 +150,17 @@
     white_channels=2,
     white_to_color_ratio=80,
 )
+FAKE_SOCKET_WITH_POWER_MONITORING = BulbType(
+    bulb_type=BulbClass.SOCKET,
+    name="ESP25_SOCKET_01",
+    features=Features(
+        color=False, color_tmp=False, effect=False, brightness=False, dual_head=False
+    ),
+    kelvin_range=KelvinRange(2700, 6500),
+    fw_version="1.26.2",
+    white_channels=2,
+    white_to_color_ratio=80,
+)
 FAKE_OLD_FIRMWARE_DIMMABLE_BULB = BulbType(
     bulb_type=BulbClass.DW,
     name=None,
@@ -197,6 +208,7 @@ async def _save_setup_callback(callback: Callable) -> None:
     )
     bulb.getMac = AsyncMock(return_value=FAKE_MAC)
     bulb.turn_on = AsyncMock()
+    bulb.get_power = AsyncMock(return_value=None)
     bulb.turn_off = AsyncMock()
     bulb.updateState = AsyncMock(return_value=FAKE_STATE)
     bulb.getSupportedScenes = AsyncMock(return_value=list(SCENES.values()))

tests/components/wiz/test_sensor.py

@@ -1,11 +1,15 @@
 """Tests for the sensor platform."""
 
+from unittest.mock import AsyncMock
+
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
 from . import (
     FAKE_DUAL_HEAD_RGBWW_BULB,
     FAKE_MAC,
+    FAKE_SOCKET_WITH_POWER_MONITORING,
+    _mocked_wizlight,
     _patch_discovery,
     _patch_wizlight,
     async_push_update,
@@ -35,3 +39,26 @@ async def test_signal_strength(hass: HomeAssistant) -> None:
 
     await async_push_update(hass, bulb, {"mac": FAKE_MAC, "rssi": -50})
     assert hass.states.get(entity_id).state == "-50"
+
+
+async def test_power_monitoring(hass: HomeAssistant) -> None:
+    """Test power monitoring."""
+    socket = _mocked_wizlight(None, None, FAKE_SOCKET_WITH_POWER_MONITORING)
+    socket.get_power = AsyncMock(return_value=5.123)
+    _, entry = await async_setup_integration(
+        hass, wizlight=socket, bulb_type=FAKE_SOCKET_WITH_POWER_MONITORING
+    )
+    entity_id = "sensor.mock_title_current_power"
+    entity_registry = er.async_get(hass)
+    reg_entry = entity_registry.async_get(entity_id)
+    assert reg_entry.unique_id == f"{FAKE_MAC}_power"
+    updated_entity = entity_registry.async_update_entity(
+        entity_id=entity_id, disabled_by=None
+    )
+    assert not updated_entity.disabled
+
+    with _patch_discovery(), _patch_wizlight(device=socket):
+        await hass.config_entries.async_reload(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert hass.states.get(entity_id).state == "5.123"