Add energy support for zwave_js meter CC entities (#53665)

* Add energy support for zwave_js meter CC entities

* shrink

* comments

* comments

* comments

* Move attributes

* Add tests

* Apply suggestions from code review

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/zwave_js/discovery.py

@@ -517,10 +517,10 @@ def get_config_parameter_discovery_schema(
         ),
         entity_registry_enabled_default=False,
     ),
-    # numeric sensors for Meter CC
+    # Meter sensors for Meter CC
     ZWaveDiscoverySchema(
         platform="sensor",
-        hint="numeric_sensor",
+        hint="meter",
         primary_value=ZWaveValueDiscoverySchema(
             command_class={
                 CommandClass.METER,

homeassistant/components/zwave_js/sensor.py

@@ -11,6 +11,7 @@
 from zwave_js_server.model.value import ConfigurationValue
 
 from homeassistant.components.sensor import (
+    ATTR_LAST_RESET,
     DEVICE_CLASS_BATTERY,
     DEVICE_CLASS_ENERGY,
     DEVICE_CLASS_ILLUMINANCE,
@@ -28,8 +29,13 @@
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_platform
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
+from homeassistant.util import dt
 
 from .const import ATTR_METER_TYPE, ATTR_VALUE, DATA_CLIENT, DOMAIN, SERVICE_RESET_METER
 from .discovery import ZwaveDiscoveryInfo
@@ -60,6 +66,8 @@ def async_add_sensor(info: ZwaveDiscoveryInfo) -> None:
             entities.append(ZWaveListSensor(config_entry, client, info))
         elif info.platform_hint == "config_parameter":
             entities.append(ZWaveConfigParameterSensor(config_entry, client, info))
+        elif info.platform_hint == "meter":
+            entities.append(ZWaveMeterSensor(config_entry, client, info))
         else:
             LOGGER.warning(
                 "Sensor not implemented for %s/%s",
@@ -128,10 +136,6 @@ def _get_device_class(self) -> str | None:
         """
         if self.info.primary_value.command_class == CommandClass.BATTERY:
             return DEVICE_CLASS_BATTERY
-        if self.info.primary_value.command_class == CommandClass.METER:
-            if self.info.primary_value.metadata.unit == "kWh":
-                return DEVICE_CLASS_ENERGY
-            return DEVICE_CLASS_POWER
         if isinstance(self.info.primary_value.property_, str):
             property_lower = self.info.primary_value.property_.lower()
             if "humidity" in property_lower:
@@ -221,14 +225,72 @@ def unit_of_measurement(self) -> str | None:
 
         return str(self.info.primary_value.metadata.unit)
 
+
+class ZWaveMeterSensor(ZWaveNumericSensor, RestoreEntity):
+    """Representation of a Z-Wave Meter CC sensor."""
+
+    _attr_state_class = STATE_CLASS_MEASUREMENT
+
+    def __init__(
+        self,
+        config_entry: ConfigEntry,
+        client: ZwaveClient,
+        info: ZwaveDiscoveryInfo,
+    ) -> None:
+        """Initialize a ZWaveNumericSensor entity."""
+        super().__init__(config_entry, client, info)
+
+        # Entity class attributes
+        self._attr_last_reset = dt.utc_from_timestamp(0)
+        self._attr_device_class = DEVICE_CLASS_POWER
+        if self.info.primary_value.metadata.unit == "kWh":
+            self._attr_device_class = DEVICE_CLASS_ENERGY
+
+    @callback
+    def async_update_last_reset(
+        self, node: ZwaveNode, endpoint: int, meter_type: int | None
+    ) -> None:
+        """Update last reset."""
+        # If the signal is not for this node or is for a different endpoint, ignore it
+        if self.info.node != node or self.info.primary_value.endpoint != endpoint:
+            return
+        # If a meter type was specified and doesn't match this entity's meter type,
+        # ignore it
+        if (
+            meter_type is not None
+            and self.info.primary_value.metadata.cc_specific.get("meterType")
+            != meter_type
+        ):
+            return
+
+        self._attr_last_reset = dt.utcnow()
+        self.async_write_ha_state()
+
+    async def async_added_to_hass(self) -> None:
+        """Call when entity is added."""
+        await super().async_added_to_hass()
+
+        # Restore the last reset time from stored state
+        restored_state = await self.async_get_last_state()
+        if restored_state and ATTR_LAST_RESET in restored_state.attributes:
+            self._attr_last_reset = dt.parse_datetime(
+                restored_state.attributes[ATTR_LAST_RESET]
+            )
+
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                f"{DOMAIN}_{SERVICE_RESET_METER}",
+                self.async_update_last_reset,
+            )
+        )
+
     async def async_reset_meter(
         self, meter_type: int | None = None, value: int | None = None
     ) -> None:
         """Reset meter(s) on device."""
         node = self.info.node
         primary_value = self.info.primary_value
-        if primary_value.command_class != CommandClass.METER:
-            raise TypeError("Reset only available for Meter sensors")
         options = {}
         if meter_type is not None:
             options["type"] = meter_type
@@ -244,6 +306,15 @@ async def async_reset_meter(
             primary_value.endpoint,
             options,
         )
+        self._attr_last_reset = dt.utcnow()
+        # Notify meters that may have been reset
+        async_dispatcher_send(
+            self.hass,
+            f"{DOMAIN}_{SERVICE_RESET_METER}",
+            node,
+            primary_value.endpoint,
+            options.get("type"),
+        )
 
 
 class ZWaveListSensor(ZwaveSensorBase):

tests/components/zwave_js/common.py

@@ -1,4 +1,6 @@
 """Provide common test tools for Z-Wave JS."""
+from datetime import datetime, timezone
+
 AIR_TEMPERATURE_SENSOR = "sensor.multisensor_6_air_temperature"
 HUMIDITY_SENSOR = "sensor.multisensor_6_humidity"
 ENERGY_SENSOR = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed_2"
@@ -29,3 +31,7 @@
     "sensor.z_wave_module_for_id_lock_150_and_101_config_parameter_door_lock_mode"
 )
 ZEN_31_ENTITY = "light.kitchen_under_cabinet_lights"
+METER_SENSOR = "sensor.smart_switch_6_electric_consumed_v"
+
+DATETIME_ZERO = datetime(1970, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
+DATETIME_LAST_RESET = datetime(2020, 1, 1, 0, 0, 0, tzinfo=timezone.utc)

tests/components/zwave_js/conftest.py

@@ -11,6 +11,11 @@
 from zwave_js_server.model.node import Node
 from zwave_js_server.version import VersionInfo
 
+from homeassistant.components.sensor import ATTR_LAST_RESET
+from homeassistant.core import State
+
+from .common import DATETIME_LAST_RESET
+
 from tests.common import MockConfigEntry, load_fixture
 
 # Add-on fixtures
@@ -835,3 +840,16 @@ def aeotec_zw164_siren_fixture(client, aeotec_zw164_siren_state):
 def firmware_file_fixture():
     """Return mock firmware file stream."""
     return io.BytesIO(bytes(10))
+
+
+@pytest.fixture(name="restore_last_reset")
+def restore_last_reset_fixture():
+    """Return mock restore last reset."""
+    state = State(
+        "sensor.test", "test", {ATTR_LAST_RESET: DATETIME_LAST_RESET.isoformat()}
+    )
+    with patch(
+        "homeassistant.components.zwave_js.sensor.ZWaveMeterSensor.async_get_last_state",
+        return_value=state,
+    ):
+        yield state

tests/components/zwave_js/test_sensor.py

@@ -1,6 +1,9 @@
 """Test the Z-Wave JS sensor platform."""
+from unittest.mock import patch
+
 from zwave_js_server.event import Event
 
+from homeassistant.components.sensor import ATTR_LAST_RESET
 from homeassistant.components.zwave_js.const import (
     ATTR_METER_TYPE,
     ATTR_VALUE,
@@ -22,10 +25,13 @@
 from .common import (
     AIR_TEMPERATURE_SENSOR,
     BASIC_SENSOR,
+    DATETIME_LAST_RESET,
+    DATETIME_ZERO,
     ENERGY_SENSOR,
     HUMIDITY_SENSOR,
     ID_LOCK_CONFIG_PARAMETER_SENSOR,
     INDICATOR_SENSOR,
+    METER_SENSOR,
     NOTIFICATION_MOTION_SENSOR,
     POWER_SENSOR,
 )
@@ -171,18 +177,30 @@ async def test_reset_meter(
     integration,
 ):
     """Test reset_meter service."""
-    SENSOR = "sensor.smart_switch_6_electric_consumed_v"
     client.async_send_command.return_value = {}
     client.async_send_command_no_wait.return_value = {}
 
-    # Test successful meter reset call
-    await hass.services.async_call(
-        DOMAIN,
-        SERVICE_RESET_METER,
-        {
-            ATTR_ENTITY_ID: SENSOR,
-        },
-        blocking=True,
+    # Validate that the sensor last reset is starting from nothing
+    assert (
+        hass.states.get(METER_SENSOR).attributes[ATTR_LAST_RESET]
+        == DATETIME_ZERO.isoformat()
+    )
+
+    # Test successful meter reset call, patching utcnow so we can make sure the last
+    # reset gets updated
+    with patch("homeassistant.util.dt.utcnow", return_value=DATETIME_LAST_RESET):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_RESET_METER,
+            {
+                ATTR_ENTITY_ID: METER_SENSOR,
+            },
+            blocking=True,
+        )
+
+    assert (
+        hass.states.get(METER_SENSOR).attributes[ATTR_LAST_RESET]
+        == DATETIME_LAST_RESET.isoformat()
     )
 
     assert len(client.async_send_command_no_wait.call_args_list) == 1
@@ -199,7 +217,7 @@ async def test_reset_meter(
         DOMAIN,
         SERVICE_RESET_METER,
         {
-            ATTR_ENTITY_ID: SENSOR,
+            ATTR_ENTITY_ID: METER_SENSOR,
             ATTR_METER_TYPE: 1,
             ATTR_VALUE: 2,
         },
@@ -214,3 +232,17 @@ async def test_reset_meter(
     assert args["args"] == [{"type": 1, "targetValue": 2}]
 
     client.async_send_command_no_wait.reset_mock()
+
+
+async def test_restore_last_reset(
+    hass,
+    client,
+    aeon_smart_switch_6,
+    restore_last_reset,
+    integration,
+):
+    """Test restoring last_reset on setup."""
+    assert (
+        hass.states.get(METER_SENSOR).attributes[ATTR_LAST_RESET]
+        == DATETIME_LAST_RESET.isoformat()
+    )