From 964af88c21e538354f0ee9c9d6eca2b0fe7c81ff Mon Sep 17 00:00:00 2001 From: Diogo Gomes Date: Fri, 2 Jun 2023 13:44:36 +0100 Subject: [PATCH] Make Riemann sum sensors restore last valid state (#93674) * keep last valid state * keep last valid state * typo * increase coverage * better error handling * debug messages * increase coverage * remove random log * don't expose last_valid_state as an attribute --- .../components/integration/sensor.py | 119 ++++++++++++++++-- tests/components/integration/test_sensor.py | 96 +++++++++++++- 2 files changed, 204 insertions(+), 11 deletions(-) diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index 7e60f2c509c..b28b426d3af 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -1,16 +1,19 @@ """Numeric integration of data coming from a source sensor over time.""" from __future__ import annotations -from decimal import Decimal, DecimalException +from dataclasses import dataclass +from decimal import Decimal, DecimalException, InvalidOperation import logging -from typing import Final +from typing import Any, Final +from typing_extensions import Self import voluptuous as vol from homeassistant.components.sensor import ( PLATFORM_SCHEMA, + RestoreSensor, SensorDeviceClass, - SensorEntity, + SensorExtraStoredData, SensorStateClass, ) from homeassistant.config_entries import ConfigEntry @@ -28,7 +31,6 @@ from homeassistant.core import Event, HomeAssistant, State, callback from homeassistant.helpers import config_validation as cv, entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.event import async_track_state_change_event -from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType from .const import ( @@ -79,6 +81,53 @@ PLATFORM_SCHEMA = vol.All( ) +@dataclass +class IntegrationSensorExtraStoredData(SensorExtraStoredData): + """Object to hold extra stored data.""" + + source_entity: str | None + last_valid_state: Decimal | None + + def as_dict(self) -> dict[str, Any]: + """Return a dict representation of the utility sensor data.""" + data = super().as_dict() + data["source_entity"] = self.source_entity + data["last_valid_state"] = ( + str(self.last_valid_state) if self.last_valid_state else None + ) + return data + + @classmethod + def from_dict(cls, restored: dict[str, Any]) -> Self | None: + """Initialize a stored sensor state from a dict.""" + extra = SensorExtraStoredData.from_dict(restored) + if extra is None: + return None + + source_entity = restored.get(ATTR_SOURCE_ID) + + try: + last_valid_state = ( + Decimal(str(restored.get("last_valid_state"))) + if restored.get("last_valid_state") + else None + ) + except InvalidOperation: + # last_period is corrupted + _LOGGER.error("Could not use last_valid_state") + return None + + if last_valid_state is None: + return None + + return cls( + extra.native_value, + extra.native_unit_of_measurement, + source_entity, + last_valid_state, + ) + + async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, @@ -129,7 +178,7 @@ async def async_setup_platform( # pylint: disable-next=hass-invalid-inheritance # needs fixing -class IntegrationSensor(RestoreEntity, SensorEntity): +class IntegrationSensor(RestoreSensor): """Representation of an integration sensor.""" _attr_state_class = SensorStateClass.TOTAL @@ -160,7 +209,8 @@ class IntegrationSensor(RestoreEntity, SensorEntity): self._unit_time = UNIT_TIME[unit_time] self._unit_time_str = unit_time self._attr_icon = "mdi:chart-histogram" - self._attr_extra_state_attributes = {ATTR_SOURCE_ID: source_entity} + self._source_entity: str = source_entity + self._last_valid_state: Decimal | None = None def _unit(self, source_unit: str) -> str: """Derive unit from the source sensor, SI prefix and time unit.""" @@ -175,10 +225,28 @@ class IntegrationSensor(RestoreEntity, SensorEntity): async def async_added_to_hass(self) -> None: """Handle entity which will be added.""" await super().async_added_to_hass() - if (state := await self.async_get_last_state()) is not None: - if state.state == STATE_UNAVAILABLE: - self._attr_available = False - elif state.state != STATE_UNKNOWN: + + if (last_sensor_data := await self.async_get_last_sensor_data()) is not None: + self._state = ( + Decimal(str(last_sensor_data.native_value)) + if last_sensor_data.native_value + else last_sensor_data.last_valid_state + ) + self._attr_native_value = last_sensor_data.native_value + self._unit_of_measurement = last_sensor_data.native_unit_of_measurement + self._last_valid_state = last_sensor_data.last_valid_state + + _LOGGER.debug( + "Restored state %s and last_valid_state %s", + self._state, + self._last_valid_state, + ) + elif (state := await self.async_get_last_state()) is not None: + # legacy to be removed on 2023.10 (we are keeping this to avoid losing data during the transition) + if state.state in [STATE_UNAVAILABLE, STATE_UNKNOWN]: + if state.state == STATE_UNAVAILABLE: + self._attr_available = False + else: try: self._state = Decimal(state.state) except (DecimalException, ValueError) as err: @@ -295,6 +363,7 @@ class IntegrationSensor(RestoreEntity, SensorEntity): self._state += integral else: self._state = integral + self._last_valid_state = self._state self.async_write_ha_state() self.async_on_remove( @@ -314,3 +383,33 @@ class IntegrationSensor(RestoreEntity, SensorEntity): def native_unit_of_measurement(self) -> str | None: """Return the unit the value is expressed in.""" return self._unit_of_measurement + + @property + def extra_state_attributes(self) -> dict[str, str] | None: + """Return the state attributes of the sensor.""" + state_attr = { + ATTR_SOURCE_ID: self._source_entity, + } + + return state_attr + + @property + def extra_restore_state_data(self) -> IntegrationSensorExtraStoredData: + """Return sensor specific state data to be restored.""" + return IntegrationSensorExtraStoredData( + self.native_value, + self.native_unit_of_measurement, + self._source_entity, + self._last_valid_state, + ) + + async def async_get_last_sensor_data( + self, + ) -> IntegrationSensorExtraStoredData | 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 IntegrationSensorExtraStoredData.from_dict( + restored_last_extra_data.as_dict() + ) diff --git a/tests/components/integration/test_sensor.py b/tests/components/integration/test_sensor.py index 515ae990deb..355d13c84d6 100644 --- a/tests/components/integration/test_sensor.py +++ b/tests/components/integration/test_sensor.py @@ -19,7 +19,7 @@ from homeassistant.core import HomeAssistant, State from homeassistant.setup import async_setup_component import homeassistant.util.dt as dt_util -from tests.common import mock_restore_cache +from tests.common import mock_restore_cache, mock_restore_cache_with_extra_data @pytest.mark.parametrize("method", ["trapezoidal", "left", "right"]) @@ -163,6 +163,100 @@ async def test_restore_state(hass: HomeAssistant) -> None: assert state.state == "100.00" assert state.attributes.get("unit_of_measurement") == UnitOfEnergy.KILO_WATT_HOUR assert state.attributes.get("device_class") == SensorDeviceClass.ENERGY + assert state.attributes.get("last_good_state") is None + + +async def test_restore_unavailable_state(hass: HomeAssistant) -> None: + """Test integration sensor state is restored correctly.""" + mock_restore_cache_with_extra_data( + hass, + [ + ( + State( + "sensor.integration", + STATE_UNAVAILABLE, + { + "device_class": SensorDeviceClass.ENERGY, + "unit_of_measurement": UnitOfEnergy.KILO_WATT_HOUR, + }, + ), + { + "native_value": None, + "native_unit_of_measurement": "kWh", + "source_entity": "sensor.power", + "last_valid_state": "100.00", + }, + ), + ], + ) + config = { + "sensor": { + "platform": "integration", + "name": "integration", + "source": "sensor.power", + "round": 2, + } + } + + assert await async_setup_component(hass, "sensor", config) + await hass.async_block_till_done() + + state = hass.states.get("sensor.integration") + assert state + assert state.state == "100.00" + + +@pytest.mark.parametrize( + "extra_attributes", + [ + { + "native_unit_of_measurement": "kWh", + "source_entity": "sensor.power", + "last_valid_state": "100.00", + }, + { + "native_value": None, + "native_unit_of_measurement": "kWh", + "source_entity": "sensor.power", + "last_valid_state": "None", + }, + ], +) +async def test_restore_unavailable_state_failed( + hass: HomeAssistant, extra_attributes +) -> None: + """Test integration sensor state is restored correctly.""" + mock_restore_cache_with_extra_data( + hass, + [ + ( + State( + "sensor.integration", + STATE_UNAVAILABLE, + { + "device_class": SensorDeviceClass.ENERGY, + "unit_of_measurement": UnitOfEnergy.KILO_WATT_HOUR, + }, + ), + extra_attributes, + ), + ], + ) + config = { + "sensor": { + "platform": "integration", + "name": "integration", + "source": "sensor.power", + "round": 2, + } + } + + assert await async_setup_component(hass, "sensor", config) + await hass.async_block_till_done() + + state = hass.states.get("sensor.integration") + assert state + assert state.state == STATE_UNAVAILABLE async def test_restore_state_failed(hass: HomeAssistant) -> None: