fix: Fixed issue with restoring saving session baseline entity and cu…

…rrent interval accumulative consumption sensor (5 minutes dev time)
BottlecapDave · Sep 28, 2024 · a7fa406 · a7fa406
1 parent 1dedab4
commit a7fa406
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/custom_components/octopus_energy/electricity/current_interval_accumulative_consumption_.py b/custom_components/octopus_energy/electricity/current_interval_accumulative_consumption_.py
@@ -1,6 +1,10 @@
 import logging
 
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 
 from homeassistant.helpers.update_coordinator import (
   CoordinatorEntity
@@ -109,7 +113,7 @@ async def async_added_to_hass(self):
     state = await self.async_get_last_state()
 
     if state is not None and self._state is None:
-      self._state = None if state.state == "unknown" else state.state
+      self._state = None if state.state in (STATE_UNAVAILABLE, STATE_UNKNOWN) else state.state
       self._attributes = dict_to_typed_dict(state.attributes)
 
       _LOGGER.debug(f'Restored OctopusEnergyCurrentElectricityIntervalAccumulativeConsumption state: {self._state}')
diff --git a/custom_components/octopus_energy/octoplus/saving_session_baseline.py b/custom_components/octopus_energy/octoplus/saving_session_baseline.py
@@ -3,6 +3,10 @@
 
 from homeassistant.core import HomeAssistant
 from homeassistant.util.dt import (utcnow)
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 
 from homeassistant.components.sensor import (
     RestoreSensor,
@@ -172,7 +176,7 @@ async def async_added_to_hass(self):
     state = await self.async_get_last_state()
 
     if state is not None:
-      self._state = None if state.state == "unknown" else state.state
+      self._state = None if state.state in (STATE_UNAVAILABLE, STATE_UNKNOWN) else state.state
       self._attributes = dict_to_typed_dict(state.attributes)
 
     _LOGGER.debug(f'Restored state: {self._state}')