Add baseline support for sensors which need to restore state after HA restart: #281

Author: tomer-w

custom_components/victron_mqtt/binary_sensor.py

@@ -2,6 +2,7 @@
 
 import logging
 from typing import Any
+from functools import cached_property
 
 from victron_mqtt import (
     Device as VictronVenusDevice,
@@ -77,7 +78,7 @@ def _on_update_task(self, value: Any) -> None:
         self._attr_is_on = new_val
         self.async_write_ha_state()
 
-    @property
+    @cached_property
     def is_on(self) -> bool:
         """Return the current state of the binary sensor."""
         assert self._attr_is_on is not None

custom_components/victron_mqtt/entity.py

@@ -1,8 +1,9 @@
 """Common code for Victron Venus integration."""
 
 from abc import abstractmethod
+from functools import cached_property
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from victron_mqtt import (
     Device as VictronVenusDevice,
@@ -13,6 +14,7 @@
 
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.const import EntityCategory, UnitOfTime
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
 
@@ -25,8 +27,17 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-class VictronBaseEntity(Entity):
-    """Implementation of a Victron Venus base entity."""
+class VictronBaseEntity(Entity if TYPE_CHECKING else object):  # type: ignore[misc]
+    """Mixin class for Victron Venus entities.
+    
+    This is a mixin class that provides common functionality for all Victron
+    entities. It should be used as the first base class in the inheritance list
+    together with a specific entity type (SensorEntity, SwitchEntity, etc.)
+    which provides the actual Entity base class.
+    
+    The TYPE_CHECKING conditional inheritance is used to satisfy type checkers
+    while avoiding MRO conflicts at runtime.
+    """
 
     def __init__(
         self,
@@ -156,7 +167,7 @@ def _map_metric_to_unit_of_measurement(
             return UnitOfTime.HOURS
         return metric.unit_of_measurement
 
-    @property
-    def device_info(self) -> DeviceInfo:
+    @cached_property
+    def device_info(self) -> DeviceInfo | None:
         """Return device information about the sensor."""
         return self._device_info

custom_components/victron_mqtt/number.py

@@ -2,6 +2,7 @@
 
 import logging
 from typing import Any
+from functools import cached_property
 
 from victron_mqtt import (
     Device as VictronVenusDevice,
@@ -88,7 +89,7 @@ def _on_update_task(self, value: Any) -> None:
         self._attr_native_value = value
         self.async_write_ha_state()
 
-    @property
+    @cached_property
     def native_value(self):
         """Return the current value."""
         return self._metric.value

custom_components/victron_mqtt/sensor.py

@@ -6,22 +6,26 @@
 """
 
 from typing import Any
+import logging
 
 from victron_mqtt import (
     Device as VictronVenusDevice,
+    FormulaMetric as VictronFormulaMetric,
     Metric as VictronVenusMetric,
     MetricKind,
 )
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import SensorEntity, SensorStateClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 
 from .entity import VictronBaseEntity
 from .hub import Hub
 
+_LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -53,8 +57,10 @@ def on_new_metric(
     hub.register_new_metric_callback(MetricKind.SENSOR, on_new_metric)
 
 
-class VictronSensor(VictronBaseEntity, SensorEntity):
+class VictronSensor(VictronBaseEntity, SensorEntity, RestoreEntity):  # type: ignore[misc]
     """Implementation of a Victron Venus sensor."""
+    
+    _baseline: float | None = None
 
     def __init__(
         self,
@@ -65,15 +71,48 @@ def __init__(
         installation_id: str,
     ) -> None:
         """Initialize the sensor."""
-        self._attr_native_value = metric.value
         super().__init__(
             device, metric, device_info, "sensor", simple_naming, installation_id
         )
 
-
     @callback
     def _on_update_task(self, value: Any) -> None:
+        if self._baseline is not None:
+            value += self._baseline
         if self._attr_native_value == value:
             return
         self._attr_native_value = value
         self.async_write_ha_state()
+
+    async def async_added_to_hass(self) -> None:
+        """Restore persistent state for FormulaMetric energy sensors."""
+
+        # Only restore for:
+        # 1. Total increasing sensors (like cumulative energy)
+        # 2. FormulaMetrics (calculated values)
+        should_restore = (
+            self._attr_state_class in [SensorStateClass.TOTAL_INCREASING, SensorStateClass.TOTAL]
+            and isinstance(self._metric, VictronFormulaMetric)
+        )
+        self._attr_native_value = self._metric.value
+        if not should_restore:
+            # Call parent to register update callbacks
+            await super().async_added_to_hass()            
+            return
+        
+        last_state = await self.async_get_last_state()
+        if last_state is not None and last_state.state is not None:
+            assert isinstance(self._attr_native_value, (int, float)), "sensor with stored baseline value must be numeric"
+            try:
+                self._baseline = float(last_state.state)
+                self._attr_native_value += self._baseline
+            except ValueError:
+                _LOGGER.warning(
+                    "Could not restore state for %s: invalid value '%s'",
+                    self.entity_id,
+                    last_state.state,
+                )
+            _LOGGER.info("Restored baseline of %.3f for %s", self._baseline, self.entity_id)
+        # Call parent to register update callbacks
+        await super().async_added_to_hass()            
+