Correct state restoring for Utility Meter sensors (#66851)

* fix merge

* backward compatability

* add status

* increase coverage

* increase further the coverage

* adds support for Decimal in SensorExtraStoredData

* more precise

* review

* don't restore broken last_reset

* increase coverage

* address review comments

* stale

* coverage increase

* Update homeassistant/components/utility_meter/sensor.py

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* catch corrupt files and respective tests

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/components/sensor/__init__.py

@@ -5,6 +5,7 @@ from collections.abc import Callable, Mapping
 from contextlib import suppress
 from dataclasses import dataclass
 from datetime import date, datetime, timedelta, timezone
+from decimal import Decimal, InvalidOperation as DecimalInvalidOperation
 import logging
 from math import floor, log10
 from typing import Any, Final, cast, final
@@ -487,17 +488,24 @@ class SensorEntity(Entity):
 class SensorExtraStoredData(ExtraStoredData):
     """Object to hold extra stored data."""
 
-    native_value: StateType | date | datetime
+    native_value: StateType | date | datetime | Decimal
     native_unit_of_measurement: str | None
 
     def as_dict(self) -> dict[str, Any]:
         """Return a dict representation of the sensor data."""
-        native_value: StateType | date | datetime | dict[str, str] = self.native_value
+        native_value: StateType | date | datetime | Decimal | dict[
+            str, str
+        ] = self.native_value
         if isinstance(native_value, (date, datetime)):
             native_value = {
                 "__type": str(type(native_value)),
                 "isoformat": native_value.isoformat(),
             }
+        if isinstance(native_value, Decimal):
+            native_value = {
+                "__type": str(type(native_value)),
+                "decimal_str": str(native_value),
+            }
         return {
             "native_value": native_value,
             "native_unit_of_measurement": self.native_unit_of_measurement,
@@ -517,12 +525,17 @@ class SensorExtraStoredData(ExtraStoredData):
                 native_value = dt_util.parse_datetime(native_value["isoformat"])
             elif type_ == "<class 'datetime.date'>":
                 native_value = dt_util.parse_date(native_value["isoformat"])
+            elif type_ == "<class 'decimal.Decimal'>":
+                native_value = Decimal(native_value["decimal_str"])
         except TypeError:
             # native_value is not a dict
             pass
         except KeyError:
             # native_value is a dict, but does not have all values
             return None
+        except DecimalInvalidOperation:
+            # native_value coulnd't be returned from decimal_str
+            return None
 
         return cls(native_value, native_unit_of_measurement)
 

homeassistant/components/utility_meter/sensor.py

@@ -1,17 +1,20 @@
 """Utility meter from sensors providing raw data."""
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 from decimal import Decimal, DecimalException, InvalidOperation
 import logging
+from typing import Any
 
 from croniter import croniter
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
     ATTR_LAST_RESET,
+    RestoreSensor,
     SensorDeviceClass,
-    SensorEntity,
+    SensorExtraStoredData,
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
@@ -32,7 +35,6 @@ from homeassistant.helpers.event import (
     async_track_point_in_time,
     async_track_state_change_event,
 )
-from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.start import async_at_start
 from homeassistant.helpers.template import is_number
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
@@ -247,7 +249,52 @@ async def async_setup_platform(
     )
 
 
-class UtilityMeterSensor(RestoreEntity, SensorEntity):
+@dataclass
+class UtilitySensorExtraStoredData(SensorExtraStoredData):
+    """Object to hold extra stored data."""
+
+    last_period: Decimal
+    last_reset: datetime | None
+    status: str
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return a dict representation of the utility sensor data."""
+        data = super().as_dict()
+        data["last_period"] = str(self.last_period)
+        if isinstance(self.last_reset, (datetime)):
+            data["last_reset"] = self.last_reset.isoformat()
+        data["status"] = self.status
+
+        return data
+
+    @classmethod
+    def from_dict(cls, restored: dict[str, Any]) -> UtilitySensorExtraStoredData | None:
+        """Initialize a stored sensor state from a dict."""
+        extra = SensorExtraStoredData.from_dict(restored)
+        if extra is None:
+            return None
+
+        try:
+            last_period: Decimal = Decimal(restored["last_period"])
+            last_reset: datetime | None = dt_util.parse_datetime(restored["last_reset"])
+            status: str = restored["status"]
+        except KeyError:
+            # restored is a dict, but does not have all values
+            return None
+        except InvalidOperation:
+            # last_period is corrupted
+            return None
+
+        return cls(
+            extra.native_value,
+            extra.native_unit_of_measurement,
+            last_period,
+            last_reset,
+            status,
+        )
+
+
+class UtilityMeterSensor(RestoreSensor):
     """Representation of an utility meter sensor."""
 
     def __init__(
@@ -422,7 +469,18 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
             )
         )
 
-        if state := await self.async_get_last_state():
+        if (last_sensor_data := await self.async_get_last_sensor_data()) is not None:
+            # new introduced in 2022.04
+            self._state = last_sensor_data.native_value
+            self._unit_of_measurement = last_sensor_data.native_unit_of_measurement
+            self._last_period = last_sensor_data.last_period
+            self._last_reset = last_sensor_data.last_reset
+            if last_sensor_data.status == COLLECTING:
+                # Null lambda to allow cancelling the collection on tariff change
+                self._collecting = lambda: None
+
+        elif state := await self.async_get_last_state():
+            # legacy to be removed on 2022.10 (we are keeping this to avoid utility_meter counter losses)
             try:
                 self._state = Decimal(state.state)
             except InvalidOperation:
@@ -445,7 +503,7 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
                     dt_util.parse_datetime(state.attributes.get(ATTR_LAST_RESET))
                 )
                 if state.attributes.get(ATTR_STATUS) == COLLECTING:
-                    # Fake cancellation function to init the meter in similar state
+                    # Null lambda to allow cancelling the collection on tariff change
                     self._collecting = lambda: None
 
         @callback
@@ -549,3 +607,23 @@ class UtilityMeterSensor(RestoreEntity, SensorEntity):
     def icon(self):
         """Return the icon to use in the frontend, if any."""
         return ICON
+
+    @property
+    def extra_restore_state_data(self) -> UtilitySensorExtraStoredData:
+        """Return sensor specific state data to be restored."""
+        return UtilitySensorExtraStoredData(
+            self.native_value,
+            self.native_unit_of_measurement,
+            self._last_period,
+            self._last_reset,
+            PAUSED if self._collecting is None else COLLECTING,
+        )
+
+    async def async_get_last_sensor_data(self) -> UtilitySensorExtraStoredData | None:
+        """Restore Utility Meter Sensor Extra Stored Data."""
+        if (restored_last_extra_data := await self.async_get_last_extra_data()) is None:
+            return None
+
+        return UtilitySensorExtraStoredData.from_dict(
+            restored_last_extra_data.as_dict()
+        )

tests/components/sensor/test_init.py

@@ -1,5 +1,6 @@
 """The test for sensor entity."""
 from datetime import date, datetime, timezone
+from decimal import Decimal
 
 import pytest
 from pytest import approx
@@ -227,10 +228,24 @@ RESTORE_DATA = {
             "isoformat": datetime(2020, 2, 8, 15, tzinfo=timezone.utc).isoformat(),
         },
     },
+    "Decimal": {
+        "native_unit_of_measurement": "°F",
+        "native_value": {
+            "__type": "<class 'decimal.Decimal'>",
+            "decimal_str": "123.4",
+        },
+    },
+    "BadDecimal": {
+        "native_unit_of_measurement": "°F",
+        "native_value": {
+            "__type": "<class 'decimal.Decimal'>",
+            "decimal_str": "123f",
+        },
+    },
 }
 
 
-# None | str | int | float | date | datetime:
+# None | str | int | float | date | datetime | Decimal:
 @pytest.mark.parametrize(
     "native_value, native_value_type, expected_extra_data, device_class",
     [
@@ -244,6 +259,7 @@ RESTORE_DATA = {
             RESTORE_DATA["datetime"],
             SensorDeviceClass.TIMESTAMP,
         ),
+        (Decimal("123.4"), dict, RESTORE_DATA["Decimal"], SensorDeviceClass.ENERGY),
     ],
 )
 async def test_restore_sensor_save_state(
@@ -294,6 +310,13 @@ async def test_restore_sensor_save_state(
             SensorDeviceClass.TIMESTAMP,
             "°F",
         ),
+        (
+            Decimal("123.4"),
+            Decimal,
+            RESTORE_DATA["Decimal"],
+            SensorDeviceClass.ENERGY,
+            "°F",
+        ),
         (None, type(None), None, None, None),
         (None, type(None), {}, None, None),
         (None, type(None), {"beer": 123}, None, None),
@@ -304,6 +327,7 @@ async def test_restore_sensor_save_state(
             None,
             None,
         ),
+        (None, type(None), RESTORE_DATA["BadDecimal"], SensorDeviceClass.ENERGY, None),
     ],
 )
 async def test_restore_sensor_restore_state(

tests/components/utility_meter/test_sensor.py

@@ -42,7 +42,11 @@ from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from tests.common import MockConfigEntry, async_fire_time_changed, mock_restore_cache
+from tests.common import (
+    MockConfigEntry,
+    async_fire_time_changed,
+    mock_restore_cache_with_extra_data,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -493,7 +497,7 @@ async def test_device_class(hass, yaml_config, config_entry_configs):
                 "utility_meter": {
                     "energy_bill": {
                         "source": "sensor.energy",
-                        "tariffs": ["onpeak", "midpeak", "offpeak"],
+                        "tariffs": ["onpeak", "midpeak", "offpeak", "superpeak"],
                     }
                 }
             },
@@ -508,7 +512,7 @@ async def test_device_class(hass, yaml_config, config_entry_configs):
                 "net_consumption": False,
                 "offset": 0,
                 "source": "sensor.energy",
-                "tariffs": ["onpeak", "midpeak", "offpeak"],
+                "tariffs": ["onpeak", "midpeak", "offpeak", "superpeak"],
             },
         ),
     ),
@@ -519,9 +523,11 @@ async def test_restore_state(hass, yaml_config, config_entry_config):
     hass.state = CoreState.not_running
 
     last_reset = "2020-12-21T00:00:00.013073+00:00"
-    mock_restore_cache(
+
+    mock_restore_cache_with_extra_data(
         hass,
         [
+            (
                 State(
                     "sensor.energy_bill_onpeak",
                     "3",
@@ -531,10 +537,36 @@ async def test_restore_state(hass, yaml_config, config_entry_config):
                         ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
                     },
                 ),
+                {
+                    "native_value": {
+                        "__type": "<class 'decimal.Decimal'>",
+                        "decimal_str": "3",
+                    },
+                    "native_unit_of_measurement": "kWh",
+                    "last_reset": last_reset,
+                    "last_period": "7",
+                    "status": "paused",
+                },
+            ),
+            (
                 State(
                     "sensor.energy_bill_midpeak",
-                "error",
+                    "5",
+                    attributes={
+                        ATTR_STATUS: PAUSED,
+                        ATTR_LAST_RESET: last_reset,
+                        ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
+                    },
                 ),
+                {
+                    "native_value": {
+                        "__type": "<class 'decimal.Decimal'>",
+                        "decimal_str": "3",
+                    },
+                    "native_unit_of_measurement": "kWh",
+                },
+            ),
+            (
                 State(
                     "sensor.energy_bill_offpeak",
                     "6",
@@ -544,6 +576,26 @@ async def test_restore_state(hass, yaml_config, config_entry_config):
                         ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
                     },
                 ),
+                {
+                    "native_value": {
+                        "__type": "<class 'decimal.Decimal'>",
+                        "decimal_str": "3f",
+                    },
+                    "native_unit_of_measurement": "kWh",
+                },
+            ),
+            (
+                State(
+                    "sensor.energy_bill_superpeak",
+                    "error",
+                    attributes={
+                        ATTR_STATUS: COLLECTING,
+                        ATTR_LAST_RESET: last_reset,
+                        ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
+                    },
+                ),
+                {},
+            ),
         ],
     )
 
@@ -569,7 +621,7 @@ async def test_restore_state(hass, yaml_config, config_entry_config):
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
     state = hass.states.get("sensor.energy_bill_midpeak")
-    assert state.state == STATE_UNKNOWN
+    assert state.state == "5"
 
     state = hass.states.get("sensor.energy_bill_offpeak")
     assert state.state == "6"
@@ -577,6 +629,9 @@ async def test_restore_state(hass, yaml_config, config_entry_config):
     assert state.attributes.get("last_reset") == last_reset
     assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ENERGY_KILO_WATT_HOUR
 
+    state = hass.states.get("sensor.energy_bill_superpeak")
+    assert state.state == STATE_UNKNOWN
+
     # utility_meter is loaded, now set sensors according to utility_meter:
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
     await hass.async_block_till_done()