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/__init__.py

@@ -712,7 +712,7 @@ async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
 
 
 _InfoT = TypeVar("_InfoT", bound=EntityInfo)
-_EntityT = TypeVar("_EntityT", bound="EsphomeBaseEntity[Any,Any]")
+_EntityT = TypeVar("_EntityT", bound="EsphomeEntity[Any,Any]")
 _StateT = TypeVar("_StateT", bound=EntityState)
 
 
@@ -850,7 +850,7 @@ class EsphomeEnumMapper(Generic[_EnumT, _ValT]):
         return self._inverse[value]
 
 
-class EsphomeBaseEntity(Entity, Generic[_InfoT, _StateT]):
+class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
     """Define a base esphome entity."""
 
     def __init__(
@@ -882,6 +882,22 @@ class EsphomeBaseEntity(Entity, Generic[_InfoT, _StateT]):
             )
         )
 
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                (
+                    f"esphome_{self._entry_id}"
+                    f"_update_{self._component_key}_{self._key}"
+                ),
+                self._on_state_update,
+            )
+        )
+
+    @callback
+    def _on_state_update(self) -> None:
+        # Behavior can be changed in child classes
+        self.async_write_ha_state()
+
     @callback
     def _on_device_update(self) -> None:
         """Update the entity state when device info has changed."""
@@ -890,7 +906,7 @@ class EsphomeBaseEntity(Entity, Generic[_InfoT, _StateT]):
             # Only update the HA state when the full state arrives
             # through the next entity state packet.
             return
-        self.async_write_ha_state()
+        self._on_state_update()
 
     @property
     def _entry_id(self) -> str:
@@ -962,23 +978,3 @@ class EsphomeBaseEntity(Entity, Generic[_InfoT, _StateT]):
     def should_poll(self) -> bool:
         """Disable polling."""
         return False
-
-
-class EsphomeEntity(EsphomeBaseEntity[_InfoT, _StateT]):
-    """Define a generic esphome entity."""
-
-    async def async_added_to_hass(self) -> None:
-        """Register callbacks."""
-
-        await super().async_added_to_hass()
-
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                (
-                    f"esphome_{self._entry_id}"
-                    f"_update_{self._component_key}_{self._key}"
-                ),
-                self.async_write_ha_state,
-            )
-        )

homeassistant/components/esphome/camera.py

@@ -10,11 +10,10 @@ from aiohttp import web
 from homeassistant.components import camera
 from homeassistant.components.camera import Camera
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import EsphomeBaseEntity, platform_async_setup_entry
+from . import EsphomeEntity, platform_async_setup_entry
 
 
 async def async_setup_entry(
@@ -32,34 +31,22 @@ async def async_setup_entry(
     )
 
 
-class EsphomeCamera(Camera, EsphomeBaseEntity[CameraInfo, CameraState]):
+class EsphomeCamera(Camera, EsphomeEntity[CameraInfo, CameraState]):
     """A camera implementation for ESPHome."""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize."""
         Camera.__init__(self)
-        EsphomeBaseEntity.__init__(self, *args, **kwargs)
+        EsphomeEntity.__init__(self, *args, **kwargs)
         self._image_cond = asyncio.Condition()
 
-    async def async_added_to_hass(self) -> None:
-        """Register callbacks."""
-
-        await super().async_added_to_hass()
-
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                (
-                    f"esphome_{self._entry_id}"
-                    f"_update_{self._component_key}_{self._key}"
-                ),
-                self._on_state_update,
-            )
-        )
-
-    async def _on_state_update(self) -> None:
+    @callback
+    def _on_state_update(self) -> None:
         """Notify listeners of new image when update arrives."""
-        self.async_write_ha_state()
+        super()._on_state_update()
+        self.hass.async_create_task(self._on_state_update_coro())
+
+    async def _on_state_update_coro(self) -> None:
         async with self._image_cond:
             self._image_cond.notify_all()
 

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."""