fix: Fixed restoration of entities to handle "unavailable" state

BottlecapDave · Jan 23, 2024 · 248e700 · 248e700
1 parent 5c85dae
commit 248e700
Show file tree

Hide file tree

Showing 46 changed files with 234 additions and 63 deletions.
diff --git a/custom_components/octopus_energy/coordinators/current_consumption.py b/custom_components/octopus_energy/coordinators/current_consumption.py
@@ -36,9 +36,10 @@ async def async_get_live_consumption(
     period_to = current_date + timedelta(days=1)
 
     try:
-      result = await client.async_get_smart_meter_consumption(device_id, period_from, period_to)
-      if result is not None:
-        return CurrentConsumptionCoordinatorResult(current_date, 1, refresh_rate_in_minutes, result)
+      data = await client.async_get_smart_meter_consumption(device_id, period_from, period_to)
+      if data is not None:
+        _LOGGER.debug(f'Retrieved current consumption data for {device_id}; period_from: {period_from}; period_to: {period_to}; length: {len(data)}; last_from: {data[-1]["start"] if len(data) > 0 else None}')
+        return CurrentConsumptionCoordinatorResult(current_date, 1, refresh_rate_in_minutes, data)
     except Exception as e:
       if isinstance(e, ApiException) == False:
         raise

diff --git a/custom_components/octopus_energy/cost_tracker/cost_tracker.py b/custom_components/octopus_energy/cost_tracker/cost_tracker.py
@@ -116,7 +116,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)
       # Make sure our attributes don't override any changed settings
       self._attributes.update(self._config)

diff --git a/custom_components/octopus_energy/cost_tracker/cost_tracker_off_peak.py b/custom_components/octopus_energy/cost_tracker/cost_tracker_off_peak.py
@@ -122,7 +122,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)
       # Make sure our attributes don't override any changed settings
       self._attributes.update(self._config)

diff --git a/custom_components/octopus_energy/cost_tracker/cost_tracker_peak.py b/custom_components/octopus_energy/cost_tracker/cost_tracker_peak.py
@@ -122,7 +122,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)
       # Make sure our attributes don't override any changed settings
       self._attributes.update(self._config)

diff --git a/custom_components/octopus_energy/electricity/current_accumulative_consumption.py b/custom_components/octopus_energy/electricity/current_accumulative_consumption.py
@@ -1,5 +1,9 @@
 import logging
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -131,7 +135,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 OctopusEnergyCurrentAccumulativeElectricityConsumption state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/current_accumulative_consumption_off_peak.py b/custom_components/octopus_energy/electricity/current_accumulative_consumption_off_peak.py
@@ -1,6 +1,10 @@
 import logging
 from datetime import datetime
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -130,7 +134,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 OctopusEnergyPreviousAccumulativeElectricityConsumptionOffPeak state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/current_accumulative_consumption_peak.py b/custom_components/octopus_energy/electricity/current_accumulative_consumption_peak.py
@@ -1,6 +1,10 @@
 import logging
 from datetime import datetime
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -130,7 +134,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 OctopusEnergyCurrentAccumulativeElectricityConsumptionPeak state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/current_accumulative_cost.py b/custom_components/octopus_energy/electricity/current_accumulative_cost.py
@@ -1,7 +1,10 @@
 import logging
 from datetime import datetime
-from custom_components.octopus_energy.coordinators.current_consumption import CurrentConsumptionCoordinatorResult
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -19,6 +22,7 @@
   calculate_electricity_consumption_and_cost,
 )
 
+from ..coordinators.current_consumption import CurrentConsumptionCoordinatorResult
 from .base import (OctopusEnergyElectricitySensor)
 from ..utils.attributes import dict_to_typed_dict
 
@@ -144,7 +148,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 OctopusEnergyCurrentAccumulativeElectricityCost state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/current_accumulative_cost_off_peak.py b/custom_components/octopus_energy/electricity/current_accumulative_cost_off_peak.py
@@ -1,6 +1,10 @@
 import logging
 from datetime import datetime
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -127,7 +131,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 OctopusEnergyCurrentAccumulativeElectricityCostOffPeak state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/current_accumulative_cost_peak.py b/custom_components/octopus_energy/electricity/current_accumulative_cost_peak.py
@@ -1,7 +1,10 @@
 import logging
 from datetime import datetime
-from custom_components.octopus_energy.coordinators.current_consumption import CurrentConsumptionCoordinatorResult
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -19,6 +22,7 @@
   calculate_electricity_consumption_and_cost,
 )
 
+from ..coordinators.current_consumption import CurrentConsumptionCoordinatorResult
 from .base import (OctopusEnergyElectricitySensor)
 from ..utils.attributes import dict_to_typed_dict
 
@@ -127,7 +131,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 OctopusEnergyCurrentAccumulativeElectricityCostPeak state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/current_consumption.py b/custom_components/octopus_energy/electricity/current_consumption.py
@@ -1,6 +1,9 @@
 import logging
-from custom_components.octopus_energy.coordinators.current_consumption import CurrentConsumptionCoordinatorResult
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.util.dt import (now)
 
@@ -16,6 +19,7 @@
     UnitOfEnergy
 )
 
+from ..coordinators.current_consumption import CurrentConsumptionCoordinatorResult
 from .base import (OctopusEnergyElectricitySensor)
 from ..utils.attributes import dict_to_typed_dict
 
@@ -112,7 +116,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)
 
       if "last_updated_timestamp" in self._attributes:

diff --git a/custom_components/octopus_energy/electricity/current_demand.py b/custom_components/octopus_energy/electricity/current_demand.py
@@ -1,9 +1,13 @@
-from custom_components.octopus_energy.coordinators.current_consumption import CurrentConsumptionCoordinatorResult
-from homeassistant.util.dt import (now)
 import logging
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
+from homeassistant.util.dt import (now)
+
 from homeassistant.helpers.update_coordinator import (
   CoordinatorEntity
 )
@@ -13,6 +17,7 @@
   SensorStateClass,
 )
 
+from ..coordinators.current_consumption import CurrentConsumptionCoordinatorResult
 from .base import (OctopusEnergyElectricitySensor)
 from ..utils.attributes import dict_to_typed_dict
 
@@ -92,7 +97,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)
 
       if "last_updated_timestamp" in self._attributes:

diff --git a/custom_components/octopus_energy/electricity/current_rate.py b/custom_components/octopus_energy/electricity/current_rate.py
@@ -1,6 +1,10 @@
 from datetime import timedelta
 import logging
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.util.dt import (utcnow)
@@ -153,6 +157,6 @@ 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, ['all_rates', 'applicable_rates'])
       _LOGGER.debug(f'Restored OctopusEnergyElectricityCurrentRate state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/next_rate.py b/custom_components/octopus_energy/electricity/next_rate.py
@@ -1,6 +1,10 @@
 from datetime import timedelta
 import logging
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.util.dt import (utcnow)
@@ -130,7 +134,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, ['all_rates', 'applicable_rates'])
 
       _LOGGER.debug(f'Restored OctopusEnergyElectricityNextRate state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/off_peak.py b/custom_components/octopus_energy/electricity/off_peak.py
@@ -1,6 +1,10 @@
 from datetime import timedelta
 import logging
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import generate_entity_id
 
@@ -80,7 +84,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" or state.state is None else state.state.lower() == 'on'
+      self._state = None if state.state in (STATE_UNAVAILABLE, STATE_UNKNOWN) or state.state is None else state.state.lower() == 'on'
       self._attributes = dict_to_typed_dict(state.attributes)
 
     if (self._state is None):

diff --git a/custom_components/octopus_energy/electricity/previous_accumulative_consumption.py b/custom_components/octopus_energy/electricity/previous_accumulative_consumption.py
@@ -1,5 +1,10 @@
 import logging
 from datetime import datetime
+
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -164,7 +169,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 OctopusEnergyPreviousAccumulativeElectricityConsumption state: {self._state}')

diff --git a/custom_components/octopus_energy/electricity/previous_accumulative_consumption_off_peak.py b/custom_components/octopus_energy/electricity/previous_accumulative_consumption_off_peak.py
@@ -1,6 +1,10 @@
 import logging
 from datetime import datetime
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -130,7 +134,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 OctopusEnergyPreviousAccumulativeElectricityConsumptionOffPeak state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/previous_accumulative_consumption_peak.py b/custom_components/octopus_energy/electricity/previous_accumulative_consumption_peak.py
@@ -1,6 +1,10 @@
 import logging
 from datetime import datetime
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant, callback
 
 from homeassistant.helpers.update_coordinator import (
@@ -130,7 +134,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 OctopusEnergyPreviousAccumulativeElectricityConsumptionPeak state: {self._state}')
diff --git a/custom_components/octopus_energy/electricity/previous_accumulative_cost.py b/custom_components/octopus_energy/electricity/previous_accumulative_cost.py
@@ -1,6 +1,10 @@
 import logging
 from datetime import datetime
 
+from homeassistant.const import (
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
 from homeassistant.core import HomeAssistant
 
 from homeassistant.helpers.update_coordinator import (
@@ -164,7 +168,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 OctopusEnergyPreviousAccumulativeElectricityCost state: {self._state}')