Skip to content

Commit

Permalink
Fix bug where area sensor doesn't update
Browse files Browse the repository at this point in the history
  • Loading branch information
wernerhp committed Nov 23, 2024
1 parent a4dd7a8 commit ebeed64
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 72 deletions.
127 changes: 62 additions & 65 deletions custom_components/load_shedding/__init__.py
Original file line number Diff line number Diff line change
@@ -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 (
Expand Down Expand Up @@ -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
Expand All @@ -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 = {}
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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 = {}
Expand All @@ -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 = []
Expand All @@ -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(
{
Expand Down Expand Up @@ -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]]):
Expand All @@ -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:
Expand All @@ -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", {})

Expand Down
12 changes: 5 additions & 7 deletions custom_components/load_shedding/sensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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] = []
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit ebeed64

Please sign in to comment.