From 03874d1b65600dd03b1de7520038b742e354e2d6 Mon Sep 17 00:00:00 2001 From: Diogo Gomes Date: Tue, 19 Apr 2022 08:01:52 +0100 Subject: [PATCH] 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 * catch corrupt files and respective tests Co-authored-by: Erik Montnemery --- homeassistant/components/sensor/__init__.py | 17 ++- .../components/utility_meter/sensor.py | 88 ++++++++++++++- tests/components/sensor/test_init.py | 26 ++++- tests/components/utility_meter/test_sensor.py | 101 ++++++++++++++---- 4 files changed, 201 insertions(+), 31 deletions(-) diff --git a/homeassistant/components/sensor/__init__.py b/homeassistant/components/sensor/__init__.py index 4fb6c83c441..c5b69cbac75 100644 --- a/homeassistant/components/sensor/__init__.py +++ b/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_ == "": native_value = dt_util.parse_date(native_value["isoformat"]) + elif type_ == "": + 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) diff --git a/homeassistant/components/utility_meter/sensor.py b/homeassistant/components/utility_meter/sensor.py index d1afc449903..8df86b3e5a8 100644 --- a/homeassistant/components/utility_meter/sensor.py +++ b/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() + ) diff --git a/tests/components/sensor/test_init.py b/tests/components/sensor/test_init.py index a0f6e6e8df3..4a3d0202c91 100644 --- a/tests/components/sensor/test_init.py +++ b/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": "", + "decimal_str": "123.4", + }, + }, + "BadDecimal": { + "native_unit_of_measurement": "°F", + "native_value": { + "__type": "", + "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( diff --git a/tests/components/utility_meter/test_sensor.py b/tests/components/utility_meter/test_sensor.py index b3ed413f249..2284d806c04 100644 --- a/tests/components/utility_meter/test_sensor.py +++ b/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,31 +523,79 @@ 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", - attributes={ - ATTR_STATUS: PAUSED, - ATTR_LAST_RESET: last_reset, - ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR, + ( + State( + "sensor.energy_bill_onpeak", + "3", + attributes={ + ATTR_STATUS: PAUSED, + ATTR_LAST_RESET: last_reset, + ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR, + }, + ), + { + "native_value": { + "__type": "", + "decimal_str": "3", + }, + "native_unit_of_measurement": "kWh", + "last_reset": last_reset, + "last_period": "7", + "status": "paused", }, ), - State( - "sensor.energy_bill_midpeak", - "error", - ), - State( - "sensor.energy_bill_offpeak", - "6", - attributes={ - ATTR_STATUS: COLLECTING, - ATTR_LAST_RESET: last_reset, - ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR, + ( + State( + "sensor.energy_bill_midpeak", + "5", + attributes={ + ATTR_STATUS: PAUSED, + ATTR_LAST_RESET: last_reset, + ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR, + }, + ), + { + "native_value": { + "__type": "", + "decimal_str": "3", + }, + "native_unit_of_measurement": "kWh", }, ), + ( + State( + "sensor.energy_bill_offpeak", + "6", + attributes={ + ATTR_STATUS: COLLECTING, + ATTR_LAST_RESET: last_reset, + ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR, + }, + ), + { + "native_value": { + "__type": "", + "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()