feat: implement retry mechanism for download failure, and better erro…

…r handling (#2234)

custom_components/powercalc/power_profile/loader/remote.py

@@ -1,3 +1,4 @@
+import asyncio
 import datetime
 import json
 import logging
@@ -22,6 +23,8 @@
 
 
 class RemoteLoader(Loader):
+    retry_timeout = 3
+
     def __init__(self, hass: HomeAssistant) -> None:
         self.hass = hass
         self.library_contents: dict = {}
@@ -75,15 +78,16 @@ async def get_model_listing(self, manufacturer: str, device_type: DeviceType | N
             if not device_type or device_type in model.get("device_type", DeviceType.LIGHT)
         }
 
-    async def load_model(self, manufacturer: str, model: str) -> tuple[dict, str] | None:
+    async def load_model(self, manufacturer: str, model: str, force_update: bool = False) -> tuple[dict, str] | None:
         model_info = self.model_infos.get(f"{manufacturer}/{model}")
         if not model_info:
             raise LibraryLoadingError("Model not found in library: %s/%s", manufacturer, model)
 
         storage_path = self.get_storage_path(manufacturer, model)
+        model_path = os.path.join(storage_path, "model.json")
 
         needs_update = False
-        path_exists = os.path.exists(storage_path)
+        path_exists = os.path.exists(model_path)
         if not path_exists:
             needs_update = True
 
@@ -93,16 +97,13 @@ async def load_model(self, manufacturer: str, model: str) -> tuple[dict, str] |
                 _LOGGER.debug("Remote profile is newer than local profile")
                 needs_update = True
 
-        if needs_update:
+        if needs_update or force_update:
             try:
-                await self.download_profile(manufacturer, model, storage_path)
-                self.set_last_update_time(time.time())
+                await self.download_with_retry(manufacturer, model, storage_path)
             except ProfileDownloadError as e:
                 if not path_exists:
                     raise e
-                _LOGGER.error("Failed to download profile, falling back to local profile")
-
-        model_path = os.path.join(storage_path, "model.json")
+                _LOGGER.debug("Failed to download profile, falling back to local profile")
 
         with open(model_path) as f:
             json_data = json.load(f)
@@ -145,6 +146,24 @@ def _get_remote_modification_time(model_info: dict) -> float:
             remote_modification_time = datetime.datetime.fromisoformat(remote_modification_time).timestamp()
         return remote_modification_time  # type: ignore
 
+    async def download_with_retry(self, manufacturer: str, model: str, storage_path: str) -> None:
+        """Download a file from a remote endpoint with retries"""
+        max_retries = 3
+        retry_count = 0
+
+        while retry_count < max_retries:
+            try:
+                await self.download_profile(manufacturer, model, storage_path)
+                break  # Break out of the loop if download is successful
+            except ProfileDownloadError as e:
+                _LOGGER.error(e, exc_info=e)
+                retry_count += 1
+                if retry_count == max_retries:
+                    raise ProfileDownloadError(f"Failed to download profile even after {max_retries} retries, falling back to local profile") from e
+
+                await asyncio.sleep(self.retry_timeout)
+                _LOGGER.warning("Failed to download profile, retrying... (Attempt %d of %d)", retry_count, max_retries)
+
     async def download_profile(self, manufacturer: str, model: str, storage_path: str) -> None:
         """
         Download the profile from Github using the Powercalc download API

tests/power_profile/loader/test_remote.py

@@ -41,6 +41,7 @@ def mock_library_json_response(mock_aioresponse: aioresponses) -> None:
 @pytest.fixture
 async def remote_loader(hass: HomeAssistant, mock_library_json_response: None) -> RemoteLoader:
     loader = RemoteLoader(hass)
+    loader.retry_timeout = 0
     await loader.initialize()
     return loader
 
@@ -135,6 +136,7 @@ async def test_download_profile_exception_unexpected_status_code(mock_aiorespons
     mock_aioresponse.get(
         f"{ENDPOINT_DOWNLOAD}/signify/LCA001",
         status=500,
+        repeat=True,
     )
 
     with pytest.raises(ProfileDownloadError):
@@ -151,24 +153,58 @@ async def test_exception_is_raised_on_connection_error(mock_aioresponse: aioresp
 
 
 async def test_exception_is_raised_on_github_resource_unavailable(mock_aioresponse: aioresponses, remote_loader: RemoteLoader) -> None:
+    manufacturer = "signify"
+    model = "LCA001"
+    storage_path = remote_loader.get_storage_path(manufacturer, model)
+    clear_storage_dir(storage_path)
+
     remote_file = {
         "path": "color_temp.csv.gz",
         "url": "https://raw.githubusercontent.com/bramstroker/homeassistant-powercalc/master/profile_library/signify/LCA001/color_temp.csv.gz",
     }
 
     mock_aioresponse.get(
-        f"{ENDPOINT_DOWNLOAD}/signify/LCA001",
+        f"{ENDPOINT_DOWNLOAD}/{manufacturer}/{model}",
         status=200,
         payload=[remote_file],
+        repeat=True,
     )
 
     mock_aioresponse.get(
         remote_file["url"],
         status=500,
+        repeat=True,
     )
 
+    remote_loader.retry_timeout = 0
     with pytest.raises(ProfileDownloadError):
-        await remote_loader.download_profile("signify", "LCA001", get_test_profile_dir("download"))
+        await remote_loader.load_model(manufacturer, model)
+
+
+async def test_eventual_success_after_download_retry(mock_aioresponse: aioresponses, remote_loader: RemoteLoader) -> None:
+    manufacturer = "signify"
+    model = "LCA001"
+    storage_path = remote_loader.get_storage_path(manufacturer, model)
+    clear_storage_dir(storage_path)
+
+    remote_file = {
+        "path": "color_temp.csv.gz",
+        "url": "https://raw.githubusercontent.com/bramstroker/homeassistant-powercalc/master/profile_library/signify/LCA001/color_temp.csv.gz",
+    }
+
+    mock_aioresponse.get(
+        f"{ENDPOINT_DOWNLOAD}/{manufacturer}/{model}",
+        status=200,
+        payload=[remote_file],
+        repeat=True,
+    )
+
+    mock_aioresponse.get(remote_file["url"], status=500)
+    mock_aioresponse.get(remote_file["url"], status=200)
+
+    await remote_loader.download_with_retry(manufacturer, model, storage_path)
+
+    assert os.path.exists(storage_path)
 
 
 @pytest.mark.parametrize(
@@ -223,7 +259,7 @@ def _mock_library_json(profile_updated_at: str) -> None:
     # Clean local directory first so we have consistent test results
     # When scenario exists_locally=True, we download the profile first, to fake the local existence
     local_storage_path = loader.get_storage_path("signify", "LCA001")
-    shutil.rmtree(local_storage_path, ignore_errors=True)
+    clear_storage_dir(local_storage_path)
 
     if exists_locally:
         loader.set_last_update_time(time.time())
@@ -256,29 +292,42 @@ async def test_fallback_to_local_profile(
     hass: HomeAssistant,
     mock_aioresponse: aioresponses,
     mock_library_json_response: None,
+    remote_loader: RemoteLoader,
 ) -> None:
-    loader = RemoteLoader(hass)
-    await loader.initialize()
-
-    local_storage_path = loader.get_storage_path("signify", "LCA001")
-    shutil.rmtree(local_storage_path, ignore_errors=True)
-    shutil.copytree(get_library_path("signify/LCA001"), local_storage_path)
+    manufacturer = "signify"
+    model = "LCA001"
+    local_storage_path = remote_loader.get_storage_path(manufacturer, model)
+    clear_storage_dir(local_storage_path)
+    shutil.copytree(get_library_path(f"{manufacturer}/{model}"), local_storage_path)
 
     mock_aioresponse.get(
-        f"{ENDPOINT_DOWNLOAD}/signify/LCA001",
+        f"{ENDPOINT_DOWNLOAD}/{manufacturer}/{model}",
         status=500,
         repeat=True,
     )
 
-    await loader.load_model("signify", "LCA001")
+    await remote_loader.load_model(manufacturer, model, force_update=True)
+
+
+async def test_profile_redownloaded_when_model_json_missing(
+    hass: HomeAssistant,
+    remote_loader: RemoteLoader,
+    mock_download_profile_endpoints: list[dict],
+) -> None:
+    local_storage_path = remote_loader.get_storage_path("signify", "LCA001")
+    shutil.rmtree(local_storage_path, ignore_errors=True)
+    os.makedirs(local_storage_path)
+
+    (__, storage_path) = await remote_loader.load_model("signify", "LCA001")
+    assert storage_path == local_storage_path
 
 
 async def test_find_model(remote_loader: RemoteLoader) -> None:
     model = await remote_loader.find_model("apple", {"HomePod (gen 2)"})
     assert model == "MQJ83"
 
 
-async def _create_loader(hass: HomeAssistant) -> RemoteLoader:
-    loader = RemoteLoader(hass)
-    await loader.initialize()
-    return loader
+def clear_storage_dir(storage_path: str) -> None:
+    if not os.path.exists(storage_path):
+        return
+    shutil.rmtree(storage_path, ignore_errors=True)

tests/power_profile/test_library.py

@@ -1,6 +1,6 @@
 import logging
 import os.path
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from homeassistant.core import HomeAssistant
@@ -159,3 +159,26 @@ async def test_linked_lut_loading(hass: HomeAssistant) -> None:
     assert profile.get_model_directory().endswith("signify/LCA006")
 
     assert os.path.exists(os.path.join(profile.get_model_directory(), "color_temp.csv.gz"))
+
+
+async def test_linked_profile_loading_failed(hass: HomeAssistant, caplog: pytest.LogCaptureFixture) -> None:
+    caplog.set_level(logging.ERROR)
+    library = await ProfileLibrary.factory(hass)
+
+    remote_loader_class = "custom_components.powercalc.power_profile.loader.remote.RemoteLoader"
+    with patch(f"{remote_loader_class}.load_model") as mock_load_model:
+        async def async_load_model_patch(manufacturer: str, __: str) -> tuple[dict, str] | None:
+            if manufacturer == "foo":
+                return None
+
+            return {
+                "manufacturer": "signify",
+                "model": "LCA001",
+                "linked_lut": "foo/bar",
+            }, ""
+
+        mock_load_model.side_effect = async_load_model_patch
+
+        await library.get_profile(ModelInfo("signify", "LCA001"))
+
+        assert "Linked model foo bar not found" in caplog.text

tests/testing_config/powercalc_profiles/linked_profile_unavailable/model.json

@@ -0,0 +1,6 @@
+{
+    "name": "Linked profile",
+    "standby_power": 0.4,
+    "calculation_strategy": "lut",
+    "linked_lut": "foo/bar"
+}