fix code and add rate-limit exception

duhow · Oct 26, 2024 · 0fc0e7b · 0fc0e7b
1 parent 6524161
commit 0fc0e7b
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 15 deletions.
diff --git a/custom_components/aigues_barcelona/api.py b/custom_components/aigues_barcelona/api.py
@@ -82,6 +82,8 @@ def _query(self, path, query=None, json=None, headers=None, method="GET"):
             raise Exception(f"Denied: {msg}")
         if resp.status_code == 400:
             raise Exception(f"Bad response: {msg}")
+        if resp.status_code == 429:
+            raise Exception(f"Rate-Limited: {msg}")
 
         return resp
 

diff --git a/custom_components/aigues_barcelona/service.py b/custom_components/aigues_barcelona/service.py
@@ -17,7 +17,7 @@ async def handle_reset_and_refresh_data(call: ServiceCall) -> None:
             _LOGGER.error("No contracts available")
             return
 
-        coordinator = contract.get("coordinator")
+        coordinator = hass.data[DOMAIN].get(contract).get("coordinator")
         if not coordinator:
             _LOGGER.error(f"Contract coordinator for {contract} not found")
             return
@@ -26,23 +26,27 @@ async def handle_reset_and_refresh_data(call: ServiceCall) -> None:
 
         # TODO: Not working - Detected unsafe call not in recorder thread
         #await clear_stored_data(hass, coordinator)
-        await fetch_historic_data(hass, coordinator)
+        await fetch_historic_data(hass, coordinator, contract)
 
     hass.services.async_register(DOMAIN, "reset_and_refresh_data", handle_reset_and_refresh_data)
     return True
 
 async def clear_stored_data(hass: HomeAssistant, coordinator) -> None:
     await coordinator._clear_statistics()
 
-async def fetch_historic_data(hass: HomeAssistant, coordinator) -> None:
+async def fetch_historic_data(hass: HomeAssistant, coordinator, contract: str = None) -> None:
     today = datetime.now()
     one_year_ago = today - timedelta(days=365)
 
     current_date = one_year_ago
     while current_date < today:
-        await coordinator._async_import_statistics(
-            await hass.async_add_executor_job(
-                coordinator._api.consumptions_week, current_date
-            )
+        consumptions = await hass.async_add_executor_job(
+            coordinator._api.consumptions_week, current_date, contract
         )
+
+        if not consumptions:
+            _LOGGER.warning(f"No data available for {current_date}")
+        else:
+            await coordinator._async_import_statistics(consumptions)
+
         current_date += timedelta(weeks=1)
diff --git a/custom_components/aigues_barcelona/services.yaml b/custom_components/aigues_barcelona/services.yaml
@@ -1,8 +1,4 @@
-# Service call documentation for Aigues de Barcelona integration
-
-reset_and_refresh_data:
-  description: "Reset all stored historic/metrics data and start getting metrics from 1 year ago, querying data weekly."
-  fields:
-    date_from:
-      description: "Optional date to start fetching metrics from (default is 1 year ago)."
-      example: "2022-01-01"
+# Service call documentation for Aigues de Barcelona integration
+
+reset_and_refresh_data:
+  description: "WARNING! Reset all stored historic/metrics data and start getting metrics from 1 year ago, querying data weekly."