Fix bug where area sensor doesn't update

Author: wernerhp

custom_components/load_shedding/__init__.py

@@ -1,12 +1,13 @@
 """The LoadShedding component."""
 from __future__ import annotations
 
-from datetime import datetime, timedelta, timezone
+from datetime import UTC, datetime, timedelta, timezone
 import logging
 from typing import Any
 
 from load_shedding.libs.sepush import SePush, SePushError
 from load_shedding.providers import Area, Stage
+import urllib3
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
@@ -186,7 +187,7 @@ def __init__(self, hass: HomeAssistant, sepush: SePush) -> None:
     async def _async_update_data(self) -> dict:
         """Retrieve latest load shedding data."""
 
-        now = datetime.now(datetime.UTC).replace(microsecond=0)
+        now = datetime.now(UTC).replace(microsecond=0)
         diff = 0
         if self.last_update is not None:
             diff = (now - self.last_update).seconds
@@ -196,6 +197,9 @@ async def _async_update_data(self) -> dict:
 
         try:
             stage = await self.async_update_stage()
+        except SePushError as err:
+            _LOGGER.error("Unable to get stage: %s", err)
+            self.data = {}
         except UpdateFailed as err:
             _LOGGER.exception("Unable to get stage: %s", err)
             self.data = {}
@@ -207,60 +211,55 @@ async def _async_update_data(self) -> dict:
 
     async def async_update_stage(self) -> dict:
         """Retrieve latest stage."""
-        now = datetime.now(datetime.UTC).replace(microsecond=0)
-        try:
-            esp = await self.hass.async_add_executor_job(self.sepush.status)
-        except SePushError as err:
-            raise UpdateFailed(err) from err
-        else:
-            data = {}
-            statuses = esp.get("status", {})
-            for idx, area in statuses.items():
-                stage = Stage(int(area.get("stage", "0")))
-                start_time = datetime.fromisoformat(area.get("stage_updated"))
+        now = datetime.now(UTC).replace(microsecond=0)
+        esp = await self.hass.async_add_executor_job(self.sepush.status)
+
+        data = {}
+        statuses = esp.get("status", {})
+        for idx, area in statuses.items():
+            stage = Stage(int(area.get("stage", "0")))
+            start_time = datetime.fromisoformat(area.get("stage_updated"))
+            start_time = start_time.replace(second=0, microsecond=0)
+            planned = [
+                {
+                    ATTR_STAGE: stage,
+                    ATTR_START_TIME: start_time.astimezone(UTC),
+                }
+            ]
+
+            next_stages = area.get("next_stages", [])
+            for i, next_stage in enumerate(next_stages):
+                # Prev
+                prev_end = datetime.fromisoformat(
+                    next_stage.get("stage_start_timestamp")
+                )
+                prev_end = prev_end.replace(second=0, microsecond=0)
+                planned[i][ATTR_END_TIME] = prev_end.astimezone(UTC)
+
+                # Next
+                stage = Stage(int(next_stage.get("stage", "0")))
+                start_time = datetime.fromisoformat(
+                    next_stage.get("stage_start_timestamp")
+                )
                 start_time = start_time.replace(second=0, microsecond=0)
-                planned = [
+                planned.append(
                     {
                         ATTR_STAGE: stage,
-                        ATTR_START_TIME: start_time.astimezone(datetime.UTC),
+                        ATTR_START_TIME: start_time.astimezone(UTC),
                     }
-                ]
-
-                next_stages = area.get("next_stages", [])
-                for i, next_stage in enumerate(next_stages):
-                    # Prev
-                    prev_end = datetime.fromisoformat(
-                        next_stage.get("stage_start_timestamp")
-                    )
-                    prev_end = prev_end.replace(second=0, microsecond=0)
-                    planned[i][ATTR_END_TIME] = prev_end.astimezone(datetime.UTC)
-
-                    # Next
-                    stage = Stage(int(next_stage.get("stage", "0")))
-                    start_time = datetime.fromisoformat(
-                        next_stage.get("stage_start_timestamp")
-                    )
-                    start_time = start_time.replace(second=0, microsecond=0)
-                    planned.append(
-                        {
-                            ATTR_STAGE: stage,
-                            ATTR_START_TIME: start_time.astimezone(datetime.UTC),
-                        }
-                    )
+                )
 
-                filtered = []
-                for stage in planned:
-                    if ATTR_END_TIME not in stage:
-                        stage[ATTR_END_TIME] = stage[ATTR_START_TIME] + timedelta(
-                            days=7
-                        )
-                    if ATTR_END_TIME in stage and stage.get(ATTR_END_TIME) >= now:
-                        filtered.append(stage)
+            filtered = []
+            for stage in planned:
+                if ATTR_END_TIME not in stage:
+                    stage[ATTR_END_TIME] = stage[ATTR_START_TIME] + timedelta(days=7)
+                if ATTR_END_TIME in stage and stage.get(ATTR_END_TIME) >= now:
+                    filtered.append(stage)
 
-                data[idx] = {
-                    ATTR_NAME: area.get("name", ""),
-                    ATTR_PLANNED: filtered,
-                }
+            data[idx] = {
+                ATTR_NAME: area.get("name", ""),
+                ATTR_PLANNED: filtered,
+            }
 
         return data
 
@@ -289,7 +288,7 @@ def add_area(self, area: Area = None) -> None:
     async def _async_update_data(self) -> dict:
         """Retrieve latest load shedding data."""
 
-        now = datetime.now(datetime.UTC).replace(microsecond=0)
+        now = datetime.now(UTC).replace(microsecond=0)
         diff = 0
         if self.last_update is not None:
             diff = (now - self.last_update).seconds
@@ -300,6 +299,9 @@ async def _async_update_data(self) -> dict:
 
         try:
             area = await self.async_update_area()
+        except SePushError as err:
+            _LOGGER.error("Unable to get area schedule: %s", err)
+            self.data = {}
         except UpdateFailed as err:
             _LOGGER.exception("Unable to get area schedule: %s", err)
             self.data = {}
@@ -315,10 +317,7 @@ async def async_update_area(self) -> dict:
         area_id_data: dict = {}
 
         for area in self.areas:
-            try:
-                esp = await self.hass.async_add_executor_job(self.sepush.area, area.id)
-            except SePushError as err:
-                raise UpdateFailed(err) from err
+            esp = await self.hass.async_add_executor_job(self.sepush.area, area.id)
 
             # Get events for area
             events = []
@@ -332,10 +331,8 @@ async def async_update_area(self) -> dict:
                     if note == str(Stage.LOAD_REDUCTION):
                         stage = Stage.LOAD_REDUCTION
 
-                start = datetime.fromisoformat(event.get("start")).astimezone(
-                    datetime.UTC
-                )
-                end = datetime.fromisoformat(event.get("end")).astimezone(datetime.UTC)
+                start = datetime.fromisoformat(event.get("start")).astimezone(UTC)
+                end = datetime.fromisoformat(event.get("end")).astimezone(UTC)
 
                 events.append(
                     {
@@ -478,7 +475,7 @@ def utc_dt(date: datetime, time: datetime) -> datetime:
         second=0,
         microsecond=0,
         tzinfo=sast,
-    ).astimezone(datetime.UTC)
+    ).astimezone(UTC)
 
 
 class LoadSheddingQuotaCoordinator(DataUpdateCoordinator[dict[str, Any]]):
@@ -494,9 +491,12 @@ def __init__(self, hass: HomeAssistant, sepush: SePush) -> None:
     async def _async_update_data(self) -> dict:
         """Retrieve latest load shedding data."""
 
-        now = datetime.now(datetime.UTC).replace(microsecond=0)
+        now = datetime.now(UTC).replace(microsecond=0)
         try:
             quota = await self.async_update_quota()
+        except SePushError as err:
+            _LOGGER.error("Unable to get quota: %s", err)
+            self.data = {}
         except UpdateFailed as err:
             _LOGGER.exception("Unable to get quota: %s", err)
         else:
@@ -507,10 +507,7 @@ async def _async_update_data(self) -> dict:
 
     async def async_update_quota(self) -> dict:
         """Retrieve latest quota."""
-        try:
-            esp = await self.hass.async_add_executor_job(self.sepush.check_allowance)
-        except SePushError as err:
-            raise UpdateFailed(err) from err
+        esp = await self.hass.async_add_executor_job(self.sepush.check_allowance)
 
         return esp.get("allowance", {})
 

custom_components/load_shedding/sensor.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import UTC, datetime, timedelta
 from typing import Any, cast
 
 from load_shedding.providers import Area, Stage
@@ -153,7 +153,7 @@ def extra_state_attributes(self) -> dict[str, list, Any]:
         if not self.data:
             return self._attr_extra_state_attributes
 
-        now = datetime.now(datetime.UTC)
+        now = datetime.now(UTC)
         # data = get_sensor_attrs(planned, planned[0].get(ATTR_STAGE, Stage.UNKNOWN))
         # data[ATTR_PLANNED] = []
         data = dict(self._attr_extra_state_attributes)
@@ -200,8 +200,6 @@ class LoadSheddingAreaSensorEntity(
 ):
     """Define a LoadShedding Area sensor entity."""
 
-    coordinator: CoordinatorEntity
-
     def __init__(self, coordinator: CoordinatorEntity, area: Area) -> None:
         """Initialize."""
         super().__init__(coordinator)
@@ -241,7 +239,7 @@ def native_value(self) -> StateType:
         if not events:
             return STATE_OFF
 
-        now = datetime.now(datetime.UTC)
+        now = datetime.now(UTC)
 
         for event in events:
             if ATTR_END_TIME in event and event.get(ATTR_END_TIME) < now:
@@ -270,7 +268,7 @@ def extra_state_attributes(self) -> dict[str, list, Any]:
         if not self.data:
             return self._attr_extra_state_attributes
 
-        now = datetime.now(datetime.UTC)
+        now = datetime.now(UTC)
         data = dict(self._attr_extra_state_attributes)
         if events := self.data.get(ATTR_FORECAST, []):
             data[ATTR_FORECAST] = []
@@ -388,7 +386,7 @@ def get_sensor_attrs(forecast: list, stage: Stage = Stage.NO_LOAD_SHEDDING) -> d
             ATTR_STAGE: stage.value,
         }
 
-    now = datetime.now(datetime.UTC)
+    now = datetime.now(UTC)
     data = dict(DEFAULT_DATA)
     data[ATTR_STAGE] = stage.value