Auto reset on value going back to 0 in ESPHome (#53592)

* ESPHome - Auto reset on value going back to 0

* Remove logging lines

* Remove useless stuff

* Move callback to sensor class
Wrap `track_change_event` in `async_on_remove`

* Convert to using internal callbacks and RestoreEntity

* Don't document fixmes?

* Review fixes

* Review fixes

Co-authored-by: Otto winter <otto@otto-winter.com>

homeassistant/components/esphome/sensor.py

@@ -1,6 +1,8 @@
 """Support for esphome sensors."""
 from __future__ import annotations
 
+from contextlib import suppress
+from datetime import datetime
 import math
 from typing import cast
 
@@ -11,6 +13,7 @@ from aioesphomeapi import (
     TextSensorInfo,
     TextSensorState,
 )
+from aioesphomeapi.model import LastResetType
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
@@ -20,9 +23,10 @@ from homeassistant.components.sensor import (
     SensorEntity,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.util import dt
 
 from . import (
@@ -71,9 +75,79 @@ _STATE_CLASSES: EsphomeEnumMapper[SensorStateClass, str | None] = EsphomeEnumMap
 )
 
 
-class EsphomeSensor(EsphomeEntity[SensorInfo, SensorState], SensorEntity):
+class EsphomeSensor(
+    EsphomeEntity[SensorInfo, SensorState], SensorEntity, RestoreEntity
+):
     """A sensor implementation for esphome."""
 
+    _old_state: float | None = None
+
+    async def async_added_to_hass(self) -> None:
+        """Register callbacks."""
+        await super().async_added_to_hass()
+
+        if self._static_info.last_reset_type != LastResetType.AUTO:
+            return
+
+        # Logic to restore old state for last_reset_type AUTO:
+        last_state = await self.async_get_last_state()
+        if last_state is None:
+            return
+
+        if "last_reset" in last_state.attributes:
+            self._attr_last_reset = dt.as_utc(
+                datetime.fromisoformat(last_state.attributes["last_reset"])
+            )
+
+        with suppress(ValueError):
+            self._old_state = float(last_state.state)
+
+    @callback
+    def _on_state_update(self) -> None:
+        """Check last_reset when new state arrives."""
+        if self._static_info.last_reset_type == LastResetType.NEVER:
+            self._attr_last_reset = dt.utc_from_timestamp(0)
+
+        if self._static_info.last_reset_type != LastResetType.AUTO:
+            super()._on_state_update()
+            return
+
+        # Last reset type AUTO logic for the last_reset property
+        # In this mode we automatically determine if an accumulator reset
+        # has taken place.
+        # We compare the last valid value (_old_state) with the new one.
+        # If the value has reset to 0 or has significantly reduced we say
+        # it has reset.
+        new_state: float | None = None
+        state = cast("str | None", self.state)
+        if state is not None:
+            with suppress(ValueError):
+                new_state = float(state)
+
+        did_reset = False
+        if new_state is None:
+            # New state is not a float - we'll detect the reset once we get valid data again
+            did_reset = False
+        elif self._old_state is None:
+            # First measurement we ever got for this sensor, always a reset
+            did_reset = True
+        elif new_state == 0:
+            # don't set reset if both old and new are 0
+            # we would already have detected the reset on the last state
+            did_reset = self._old_state != 0
+        elif new_state < self._old_state:
+            did_reset = True
+
+        # Set last_reset to now if we detected a reset
+        if did_reset:
+            self._attr_last_reset = dt.utcnow()
+
+        if new_state is not None:
+            # Only write to old_state if the new one contains actual data
+            self._old_state = new_state
+
+        super()._on_state_update()
+
     @property
     def icon(self) -> str | None:
         """Return the icon."""