Make RestoreSensor more resilience

Author: tomer-w

custom_components/victron_mqtt/sensor.py

@@ -131,9 +131,9 @@ def _normalize_value(value: Any) -> Any:
     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)
+        # Only restore for cumulative FormulaMetric sensors (TOTAL / TOTAL_INCREASING).
+        # These metrics start from 0 on each HA restart, so we restore the
+        # previous accumulated value as a baseline and add new increments on top.
         should_restore = self.state_class in [
             SensorStateClass.TOTAL_INCREASING,
             SensorStateClass.TOTAL,
@@ -145,23 +145,29 @@ async def async_added_to_hass(self) -> None:
             return
 
         last_state = await self.async_get_last_state()
-        if last_state is None or last_state.state in (None, "unknown"):
+        if last_state is None or last_state.state in (None, "unknown", "unavailable"):
             await super().async_added_to_hass()
             _LOGGER.info(
-                    "Baseline is missing. Probably first load for %s", self.entity_id
-                )
+                "Baseline is missing. Probably first load for %s", self.entity_id
+            )
+            return
+
+        if not isinstance(self._attr_native_value, int | float):
+            _LOGGER.warning(
+                "Cannot restore baseline for %s: current value is %r (expected numeric)",
+                self.entity_id,
+                self._attr_native_value,
+            )
+            await super().async_added_to_hass()
             return
 
-        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
             _LOGGER.info(
                 "Restored baseline of %.3f for %s", self._baseline, self.entity_id
             )
-        except ValueError:
+        except (ValueError, TypeError):
             _LOGGER.warning(
                 "Could not restore state for %s: invalid value '%s' (type: %s)",
                 self.entity_id,

tests/test_hub.py

@@ -536,11 +536,11 @@ async def test_sensor_with_baseline(
     snapshot: SnapshotAssertion,
     init_integration,
 ) -> None:
-    """Test that baseline is NOT restored for TOTAL state_class (only TOTAL_INCREASING).
-    
-    The current implementation only restores baseline for sensors with
-    SensorStateClass.TOTAL_INCREASING. FormulaMetric sensors with CUMULATIVE
-    nature get mapped to SensorStateClass.TOTAL, so baseline is not restored.
+    """Test that baseline IS restored for FormulaMetric sensors with TOTAL state_class.
+
+    FormulaMetric energy sensors start from 0 on each HA restart. The baseline
+    mechanism restores the previous accumulated value so the sensor continues
+    from where it left off, preventing spikes in the Energy Dashboard.
     """
     victron_hub, mock_config_entry = init_integration
     # Mock time.monotonic() to return a fixed time
@@ -573,12 +573,11 @@ async def test_sensor_with_baseline(
     energy_entities = [e for e in entities if "energy" in e.entity_id]
     assert len(energy_entities) > 0
 
-    # Since state_class is TOTAL (not TOTAL_INCREASING), baseline is NOT restored
-    # The value should be 0.0 (the FormulaMetric's initial value)
+    # Since state_class is TOTAL and this is a FormulaMetric, baseline IS restored
+    # The value = baseline (1000.0) + formula accumulated energy (0.004)
     state = hass.states.get(energy_entities[0].entity_id)
     assert state is not None
-    # Value should be the metric's initial value, NOT the baseline
-    assert float(state.state) == 1000.004  # Not 0.004
+    assert float(state.state) == 1000.004
     # Should have created two entity
     assert len(entities) == 2
     assert entities == snapshot