Support custom upgrade for custom repo in WLED

homeassistant/components/wled/coordinator.py

@@ -183,27 +183,44 @@ async def _async_update_data(self) -> WLEDDevice:
         return device
 
 
-class WLEDReleasesDataUpdateCoordinator(DataUpdateCoordinator[Releases]):
+class WLEDReleasesDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Releases]]):
     """Class to manage fetching WLED releases."""
 
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize global WLED releases updater."""
-        self.wled = WLEDReleases(session=async_get_clientsession(hass))
         super().__init__(
             hass,
             LOGGER,
             config_entry=None,
             name=DOMAIN,
             update_interval=RELEASES_SCAN_INTERVAL,
         )
+        self.repos: dict[str, str] = {}
 
-    async def _async_update_data(self) -> Releases:
+    async def register_repo(self, entry_id: str, repo: str) -> None:
+        """Register a repo to fetch releases for."""
+        if self.repos.get(entry_id) != repo:
+            self.repos[entry_id] = repo.lower()
+            await self.async_request_refresh()
+
+    def unregister_repo(self, entry_id: str) -> None:
+        """Unregister a repo."""
+        del self.repos[entry_id]
+
+    async def _async_update_data(self) -> dict[str, Releases]:
         """Fetch release data from WLED."""
-        try:
-            return await self.wled.releases()
-        except WLEDError as error:
-            raise UpdateFailed(
-                translation_domain=DOMAIN,
-                translation_key="invalid_response_github_error",
-                translation_placeholders={"error": str(error)},
-            ) from error
+        result = {}
+        repos = set(self.repos.values())
+        for repo in repos:
+            try:
+                update = await WLEDReleases(
+                    repo=repo,
+                    session=async_get_clientsession(self.hass),
+                ).releases()
+            except WLEDError as error:
+                self.logger.warning(
+                    "Error fetching releases for repo %s: %s", repo, error
+                )
+            else:
+                result[repo] = update
+        return result

homeassistant/components/wled/manifest.json

@@ -7,6 +7,8 @@
   "integration_type": "device",
   "iot_class": "local_push",
   "quality_scale": "platinum",
-  "requirements": ["wled==0.23.0"],
+  "requirements": [
+    "wled @ git+https://github.com/LordMike/python-wled@codex/use-device-repo-upgrade"
+  ],
   "zeroconf": ["_wled._tcp.local."]
 }

homeassistant/components/wled/update.py

@@ -2,12 +2,15 @@
 
 from typing import Any, cast
 
+from wled import Releases
+from wled.const import DEFAULT_REPO
+
 from homeassistant.components.update import (
     UpdateDeviceClass,
     UpdateEntity,
     UpdateEntityFeature,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from . import WLED_KEY
@@ -61,11 +64,47 @@
                 self._handle_coordinator_update
             )
         )
+        await self.releases_coordinator.register_repo(
+            entry_id=self.coordinator.config_entry.entry_id,
+            repo=self._repo,
+        )
+
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Handle updated data from the coordinator."""
+        self.hass.async_create_task(
+            self.releases_coordinator.register_repo(
+                entry_id=self.coordinator.config_entry.entry_id,
+                repo=self._repo,
+            )
+        )
+        super()._handle_coordinator_update()
+
+    async def async_will_remove_from_hass(self) -> None:
+        """When removed from hass."""
+        self.releases_coordinator.unregister_repo(
+            entry_id=self.coordinator.config_entry.entry_id
+        )
+        return await super().async_will_remove_from_hass()
+
+    @property
+    def _repo(self) -> str:
+        """Return the repo to fetch releases for."""
+        return self.coordinator.data.info.repo or DEFAULT_REPO
+
+    @property
+    def _release_info(self) -> Releases | None:
+        """Return the release info of the latest version available."""
+        return self.releases_coordinator.data.get(self._repo)
 
     @property
     def available(self) -> bool:
         """Return if entity is available."""
-        return super().available and self.releases_coordinator.last_update_success
+        return (
+            super().available
+            and self.releases_coordinator.last_update_success
+            and self._release_info is not None
+        )
 
     @property
     def installed_version(self) -> str | None:
@@ -77,20 +116,23 @@
     @property
     def latest_version(self) -> str | None:
         """Latest version available for install."""
+        if (release_info := self._release_info) is None:
+            return None
+
         # If we already run a pre-release, we consider being on the beta channel.
         # Offer beta version upgrade, unless stable is newer
         if (
-            (beta := self.releases_coordinator.data.beta) is not None
+            (beta := release_info.beta) is not None
             and (current := self.coordinator.data.info.version) is not None
             and (current.alpha or current.beta or current.release_candidate)
             and (
-                (stable := self.releases_coordinator.data.stable) is None
+                (stable := release_info.stable) is None
                 or (stable is not None and stable < beta and current > stable)
             )
         ):
             return str(beta)
 
-        if (stable := self.releases_coordinator.data.stable) is not None:
+        if (stable := release_info.stable) is not None:
             return str(stable)
 
         return None
@@ -100,7 +142,7 @@
         """URL to the full release notes of the latest version available."""
         if (version := self.latest_version) is None:
             return None
-        return f"https://github.com/wled/WLED/releases/tag/v{version}"
+        return f"https://github.com/{self._repo}/releases/tag/v{version}"
 
     @wled_exception_handler
     async def async_install(
@@ -110,7 +152,7 @@
         if version is None:
             # We cast here, as we know that the latest_version is a string.
             version = cast(str, self.latest_version)
-        await self.coordinator.wled.upgrade(version=version)
+        await self.coordinator.wled.upgrade(version=version, repo=self._repo)
         await self.coordinator.async_refresh()
 
     async def async_update(self) -> None: