Use device-reported repo for upgrades

Author: LordMike

src/wled/models.py

@@ -15,6 +15,7 @@
 
 from .const import (
     CUSTOM_PALETTE_ID_CHANGE_VERSION,
+    DEFAULT_REPO,
     MIN_REQUIRED_VERSION,
     LightCapability,
     LiveDataOverride,
@@ -514,6 +515,9 @@ class Info(BaseModel):  # pylint: disable=too-many-instance-attributes
     product: str = "DIY Light"
     """The product name. Always FOSS for standard installations."""
 
+    repo: str = DEFAULT_REPO
+    """GitHub repository in 'owner/repository' format."""
+
     release: str | None = None
     """The release name of the firmware build.
 

src/wled/wled.py

@@ -636,14 +636,15 @@ async def upgrade(  # noqa: PLR0912
         self,
         *,
         version: str | AwesomeVersion,
-        repo: str = DEFAULT_REPO,
+        repo: str | None = None,
     ) -> None:
         """Upgrade WLED device to the specified version.
 
         Args:
         ----
             version: The version to upgrade to.
-            repo: GitHub repository to download firmware from.
+            repo: GitHub repository to download firmware from. If not specified,
+                the repository reported by the device firmware is used.
 
         Raises:
         ------
@@ -682,6 +683,9 @@ async def upgrade(  # noqa: PLR0912
             msg = "Device already running the requested version"
             raise WLEDUpgradeError(msg)
 
+        if repo is None:
+            repo = self._device.info.repo
+
         # Determine if this is an Ethernet board
         ethernet = ""
         if (

tests/__snapshots__/test_models.ambr

@@ -759,6 +759,7 @@
       palette_count=75,
       product='FOSS',
       release='ESP32',
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=79769),
@@ -1994,6 +1995,7 @@
       palette_count=71,
       product='FOSS',
       release=None,
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=966),
@@ -3258,6 +3260,7 @@
       palette_count=71,
       product='FOSS',
       release=None,
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=461),
@@ -4469,6 +4472,7 @@
       palette_count=71,
       product='FOSS',
       release=None,
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=461),
@@ -5680,6 +5684,7 @@
       palette_count=71,
       product='FOSS',
       release=None,
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=461),
@@ -6894,6 +6899,7 @@
       palette_count=71,
       product='FOSS',
       release=None,
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=321),
@@ -8105,6 +8111,7 @@
       palette_count=71,
       product='FOSS',
       release=None,
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=72),
@@ -9320,6 +9327,7 @@
       palette_count=71,
       product='FOSS',
       release='ESP32',
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=12),
@@ -10535,6 +10543,7 @@
       palette_count=71,
       product='FOSS',
       release='ESP32',
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=129),
@@ -11750,6 +11759,7 @@
       palette_count=71,
       product='FOSS',
       release='ESP32',
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=119),
@@ -12965,6 +12975,7 @@
       palette_count=71,
       product='FOSS',
       release='ESP32',
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=47),
@@ -14180,6 +14191,7 @@
       palette_count=71,
       product='FOSS',
       release='ESP32',
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=40),
@@ -15539,6 +15551,7 @@
       palette_count=72,
       product='FOSS',
       release='ESP32',
+      repo='wled/WLED',
       sync_toggle_receive=False,
       udp_port=21324,
       uptime=datetime.timedelta(seconds=75),

tests/test_models.py

@@ -10,6 +10,7 @@
 from syrupy.assertion import SnapshotAssertion
 
 from wled import Device, Playlist, Preset, Releases
+from wled.const import DEFAULT_REPO
 from wled.exceptions import WLEDUnsupportedVersionError
 from wled.models import (
     AwesomeVersionSerializationStrategy,
@@ -315,6 +316,18 @@ def test_info_version_deserialized() -> None:
     assert str(info.version) == "0.14.0"
 
 
+def test_info_repo_defaults_to_default_repo() -> None:
+    """Test repo defaults to DEFAULT_REPO when not present."""
+    info = Info.from_dict(_base_info())
+    assert info.repo == DEFAULT_REPO
+
+
+def test_info_repo_uses_device_value_when_present() -> None:
+    """Test repo is deserialized from the device response."""
+    info = Info.from_dict(_base_info(repo="MoonModules/WLED"))
+    assert info.repo == "MoonModules/WLED"
+
+
 # =========================================================================
 # State model
 # =========================================================================

tests/test_wled.py

@@ -12,7 +12,7 @@
 from yarl import URL
 
 from wled import WLED, Device, Releases
-from wled.const import LiveDataOverride
+from wled.const import DEFAULT_REPO, LiveDataOverride
 from wled.exceptions import (
     WLEDConnectionClosedError,
     WLEDConnectionError,
@@ -1260,6 +1260,45 @@ async def test_upgrade_success(responses: aioresponses, wled: WLED) -> None:
     await wled.upgrade(version="0.15.0")
 
 
+@pytest.mark.parametrize(
+    ("repo", "download_repo"),
+    [
+        (None, "MoonModules/WLED"),
+        (DEFAULT_REPO, DEFAULT_REPO),
+    ],
+    ids=["device_repo", "explicit_repo"],
+)
+async def test_upgrade_repo_selection(
+    responses: aioresponses,
+    wled: WLED,
+    repo: str | None,
+    download_repo: str,
+) -> None:
+    """Test upgrade selects the expected firmware repository."""
+    wled_data = load_fixture_json("wled")
+    wled_data["info"]["arch"] = "esp32"
+    wled_data["info"]["ver"] = "0.14.0"
+    wled_data["info"]["repo"] = "MoonModules/WLED"
+    mock_json_and_presets(responses, wled_data)
+    await wled.update()
+    responses.get(
+        f"https://github.com/{download_repo}/releases/download/v0.15.0/"
+        "WLED_0.15.0_ESP32.bin",
+        status=200,
+        body=b"fake firmware",
+    )
+    responses.post(
+        "http://example.com/update",
+        status=200,
+        body="OK",
+        content_type="text/plain",
+    )
+    if repo is None:
+        await wled.upgrade(version="0.15.0")
+    else:
+        await wled.upgrade(version="0.15.0", repo=repo)
+
+
 async def test_upgrade_ethernet_board(responses: aioresponses, wled: WLED) -> None:
     """Test upgrade with Ethernet board (empty bssid)."""
     await prepare_wled_for_upgrade(responses, wled, wifi_bssid="")