fix: stabilize forecast usage/state handling and coordinator shutdown (#328)

- add fallback logic for forecasted usage sensor when future consumption data is missing

- fix coordinator lifecycle cleanup and timezone handling in statistics fetching

- cap stale statistics backfill window and improve forecast failure behavior

- add typecheck script and missing translation label

Author: GuyKh

AGENTS.md

@@ -175,6 +175,13 @@ When iec-api updates:
 
 ## Validation
 
+**After every code change, run:**
+```bash
+ruff check .
+ruff format .
+./scripts/typecheck   # Run mypy type checker
+```
+
 **Before committing, run:**
 ```bash
 ./scripts/lint        # Auto-format and fix linting issues

custom_components/iec/coordinator.py

@@ -91,6 +91,7 @@ def __init__(
         super().__init__(
             hass,
             _LOGGER,
+            config_entry=config_entry,
             name="Iec",
             # Data is updated daily on IEC.
             # Refresh every 1h to be at most 5h behind.
@@ -127,10 +128,14 @@ def _dummy_listener() -> None:
         # Needed when the _async_update_data below returns {} for utilities that don't provide
         # forecast, which results to no sensors added, no registered listeners, and thus
         # _async_update_data not periodically getting called which is needed for _insert_statistics.
-        self.async_add_listener(_dummy_listener)
+        self._dummy_listener_unsub = self.async_add_listener(_dummy_listener)
 
     async def async_unload(self):
         """Unload the coordinator, cancel any pending tasks."""
+        if self._dummy_listener_unsub is not None:
+            self._dummy_listener_unsub()
+            self._dummy_listener_unsub = None
+        await self.async_shutdown()
         _LOGGER.info("Coordinator unloaded successfully.")
 
     async def _get_devices_by_contract_id(self, contract_id) -> list[Device]:
@@ -876,13 +881,18 @@ async def _insert_statistics(self, contract_id: int, is_smart_meter: bool) -> No
                     ReadingResolution.MONTHLY,
                 )
 
-                if readings and readings.meter_start_date:
+                if (
+                    readings
+                    and readings.meter_list
+                    and readings.meter_list[0].meter_start_date
+                ):
                     # Fetching the last reading from either the installation date or a month ago
                     month_ago_time = max(
                         month_ago_time,
                         localize_datetime(
                             datetime.combine(
-                                readings.meter_start_date, datetime.min.time()
+                                readings.meter_list[0].meter_start_date,
+                                datetime.min.time(),
                             )
                         ),
                     )
@@ -907,7 +917,7 @@ async def _insert_statistics(self, contract_id: int, is_smart_meter: bool) -> No
             else:
                 last_stat_time = last_stat[consumption_statistic_id][0]["start"]
                 # API returns daily data, so need to increase the start date by 4 hrs to get the next day
-                from_date = datetime.fromtimestamp(last_stat_time)
+                from_date = localize_datetime(datetime.fromtimestamp(last_stat_time))
                 _LOGGER.debug(
                     f"[IEC Statistics] Last statistics are from {from_date.strftime('%Y-%m-%d %H:%M:%S')}"
                 )
@@ -923,6 +933,16 @@ async def _insert_statistics(self, contract_id: int, is_smart_meter: bool) -> No
                         hour=1, minute=0, second=0, microsecond=0
                     )
 
+                min_from_date = (localized_today - timedelta(days=30)).replace(
+                    hour=1, minute=0, second=0, microsecond=0
+                )
+                if from_date < min_from_date:
+                    _LOGGER.debug(
+                        "[IEC Statistics] Last statistics are too old, limiting fetch window to %s",
+                        min_from_date.strftime("%Y-%m-%d %H:%M:%S"),
+                    )
+                    from_date = min_from_date
+
                 _LOGGER.debug(
                     f"[IEC Statistics] Fetching consumption from {from_date.strftime('%Y-%m-%d %H:%M:%S')}"
                 )
@@ -949,7 +969,7 @@ async def _insert_statistics(self, contract_id: int, is_smart_meter: bool) -> No
                 continue
 
             last_stat_hour = (
-                datetime.fromtimestamp(last_stat_time)
+                localize_datetime(datetime.fromtimestamp(last_stat_time))
                 if last_stat_time
                 else readings.meter_list[0].period_consumptions[0].interval
             )
@@ -1214,10 +1234,10 @@ def _calculate_estimated_bill(
                 )
             else:
                 _LOGGER.warn(
-                    f"Failed to calculate Future Consumption, Assuming last meter read \
-                    ({last_meter_read}) as full consumption"
+                    f"Failed to calculate Future Consumption for meter {meter_id} "
+                    f"(missing total_import), defaulting forecasted consumption to 0"
                 )
-                future_consumption = last_meter_read
+                future_consumption = 0
 
         kva_price = power_size * kva_tariff / 365
 

custom_components/iec/sensor.py

@@ -264,14 +264,27 @@ def _get_reading_by_date(
         state_class=SensorStateClass.TOTAL_INCREASING,
         suggested_display_precision=3,
         value_fn=lambda data: (
-            (
-                data[FUTURE_CONSUMPTIONS_DICT_NAME][
+            data[FUTURE_CONSUMPTIONS_DICT_NAME][
+                data[ATTRIBUTES_DICT_NAME][METER_ID_ATTR_NAME]
+            ].total_import
+            if (
+                data[FUTURE_CONSUMPTIONS_DICT_NAME]
+                and data[FUTURE_CONSUMPTIONS_DICT_NAME].get(
+                    data[ATTRIBUTES_DICT_NAME][METER_ID_ATTR_NAME]
+                )
+                and data[FUTURE_CONSUMPTIONS_DICT_NAME][
                     data[ATTRIBUTES_DICT_NAME][METER_ID_ATTR_NAME]
                 ].total_import
-                or 0
+                is not None
+            )
+            else (
+                data[INVOICE_DICT_NAME].meter_readings[0].reading
+                if (
+                    data[INVOICE_DICT_NAME] != EMPTY_INVOICE
+                    and data[INVOICE_DICT_NAME].meter_readings
+                )
+                else None
             )
-            if (data[FUTURE_CONSUMPTIONS_DICT_NAME])
-            else None
         ),
     ),
 )

custom_components/iec/translations/en.json

@@ -12,6 +12,9 @@
             "access_token_issued_at": {
                 "name": "Access Token Issued At"
             },
+            "elec_forecasted_usage": {
+                "name": "Next electric bill forecasted usage {multi_contract}"
+            },
             "elec_forecasted_cost": {
                 "name": "Next electric bill forecasted cost {multi_contract}",
                 "state_attributes": {

scripts/typecheck

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+set -e
+
+cd "$(dirname "$0")/.."
+
+python3 -m mypy \
+    --cache-dir /tmp/mypy_cache_iec \
+    --namespace-packages \
+    --explicit-package-bases \
+    --disable-error-code=import-untyped \
+    custom_components/iec