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>
pull/70171/head^2
Diogo Gomes 2022-04-19 08:01:52 +01:00 committed by GitHub
parent 9dfd37c60b
commit 03874d1b65
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 201 additions and 31 deletions

View File

@ -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)

View File

@ -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()
)

View File

@ -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(

View File

@ -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": "<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",
),
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": "<class 'decimal.Decimal'>",
"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": "<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()