From e2d2deb192f580d0fac93ce925c1caa0416169f9 Mon Sep 17 00:00:00 2001
From: Will Tatam <will@netmindz.net>
Date: Sun, 8 Dec 2024 16:28:43 +0000
Subject: [PATCH] Replace hard coded repo with parameter

---
 examples/upgrade.py |  7 +++++--
 src/wled/const.py   |  1 +
 src/wled/models.py  |  4 ++++
 src/wled/wled.py    | 18 ++++++++++++++----
 4 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/examples/upgrade.py b/examples/upgrade.py
index 095050f3..a5d37fa3 100644
--- a/examples/upgrade.py
+++ b/examples/upgrade.py
@@ -2,12 +2,14 @@
 """Asynchronous Python client for WLED."""
 
 import asyncio
+import sys
 
 from wled import WLED, WLEDReleases
 
 
 async def main() -> None:
     """Show example on upgrade your WLED device."""
+#    async with WLEDReleases("MoonModules/WLED") as releases:
     async with WLEDReleases() as releases:
         latest = await releases.releases()
         print(f"Latest stable version: {latest.stable}")
@@ -17,12 +19,13 @@ async def main() -> None:
         print("No stable version found")
         return
 
-    async with WLED("10.10.11.54") as led:
+    async with WLED(sys.argv[1]) as led:
         device = await led.update()
         print(f"Current version: {device.info.version}")
+        print(f"Current release: {device.info.release}")
 
         print("Upgrading WLED....")
-        await led.upgrade(version=latest.stable)
+        await led.upgrade(version=latest.stable,repo=latest.repo) # stable not default option
 
         print("Waiting for WLED to come back....")
         await asyncio.sleep(5)
diff --git a/src/wled/const.py b/src/wled/const.py
index 3bd57ba8..c60448ca 100644
--- a/src/wled/const.py
+++ b/src/wled/const.py
@@ -6,6 +6,7 @@
 
 MIN_REQUIRED_VERSION = AwesomeVersion("0.14.0")
 
+WLED_DEFAULT_REPO = "Aircoookie/WLED"
 
 class LightCapability(IntFlag):
     """Enumeration representing the capabilities of a light in WLED."""
diff --git a/src/wled/models.py b/src/wled/models.py
index 7acea0f2..e7286d91 100644
--- a/src/wled/models.py
+++ b/src/wled/models.py
@@ -478,6 +478,9 @@ class Info(BaseModel):  # pylint: disable=too-many-instance-attributes
     )
     """Version of the WLED software."""
 
+    release: str | None = field(metadata=field_options(alias="release"))
+    """Release info of the WLED software."""
+
     websocket: int | None = field(default=None, metadata=field_options(alias="ws"))
     """
     Number of currently connected WebSockets clients.
@@ -847,3 +850,4 @@ class Releases(BaseModel):
 
     beta: AwesomeVersion | None
     stable: AwesomeVersion | None
+    repo: str
diff --git a/src/wled/wled.py b/src/wled/wled.py
index 1e442bc4..ae6715ef 100644
--- a/src/wled/wled.py
+++ b/src/wled/wled.py
@@ -12,6 +12,8 @@
 import orjson
 from yarl import URL
 
+from wled.const import WLED_DEFAULT_REPO
+
 from .exceptions import (
     WLEDConnectionClosedError,
     WLEDConnectionError,
@@ -572,7 +574,7 @@ async def nightlight(
         nightlight = {k: v for k, v in nightlight.items() if v is not None}
         await self.request("/json/state", method="POST", data={"nl": nightlight})
 
-    async def upgrade(self, *, version: str | AwesomeVersion) -> None:
+    async def upgrade(self, *, version: str | AwesomeVersion, repo) -> None:
         """Upgrades WLED device to the specified version.
 
         Args:
@@ -637,7 +639,7 @@ async def upgrade(self, *, version: str | AwesomeVersion) -> None:
         architecture = self._device.info.architecture.upper()
         update_file = f"WLED_{version}_{architecture}{ethernet}.bin{gzip}"
         download_url = (
-            "https://github.com/Aircoookie/WLED/releases/download"
+            f"https://github.com/{repo}/releases/download"
             f"/v{version}/{update_file}"
         )
 
@@ -656,7 +658,7 @@ async def upgrade(self, *, version: str | AwesomeVersion) -> None:
             raise WLEDConnectionTimeoutError(msg) from exception
         except aiohttp.ClientResponseError as exception:
             if exception.status == 404:
-                msg = f"Requested WLED version '{version}' does not exists"
+                msg = f"Requested WLED bin {update_file} for version '{version}' does not exists"
                 raise WLEDUpgradeError(msg) from exception
             msg = (
                 f"Could not download requested WLED version '{version}'"
@@ -707,10 +709,16 @@ class WLEDReleases:
 
     request_timeout: float = 8.0
     session: aiohttp.client.ClientSession | None = None
+    repo: str = WLED_DEFAULT_REPO
+
 
     _client: aiohttp.ClientWebSocketResponse | None = None
     _close_session: bool = False
 
+    def __init__(self, repo=None):
+        if repo is not None:
+            self.repo = repo
+
     @backoff.on_exception(backoff.expo, WLEDConnectionError, max_tries=3, logger=None)
     async def releases(self) -> Releases:
         """Fetch WLED version information from GitHub.
@@ -733,10 +741,11 @@ async def releases(self) -> Releases:
             self.session = aiohttp.ClientSession()
             self._close_session = True
 
+        print(f"https://api.github.com/repos/{self.repo}/releases")
         try:
             async with asyncio.timeout(self.request_timeout):
                 response = await self.session.get(
-                    "https://api.github.com/repos/Aircoookie/WLED/releases",
+                    f"https://api.github.com/repos/{self.repo}/releases",
                     headers={"Accept": "application/json"},
                 )
         except asyncio.TimeoutError as exception:
@@ -781,6 +790,7 @@ async def releases(self) -> Releases:
         return Releases(
             beta=version_latest_beta,
             stable=version_latest,
+            repo=self.repo,
         )
 
     async def close(self) -> None: