Add possibility to upgrade WLED device (#465)

Author: frenck

.flake8

@@ -1,4 +1,4 @@
 [flake8]
 max-line-length = 88
-ignore = D202,W503
+ignore = D202,E501,W503
 per-file-ignores = tests/*:DAR,S101

examples/upgrade.py

@@ -0,0 +1,20 @@
+# pylint: disable=W0621
+"""Asynchronous Python client for WLED."""
+
+import asyncio
+
+from wled import WLED
+
+
+async def main():
+    """Show example on upgrade your WLED device."""
+    async with WLED("10.10.11.54") as led:
+        device = await led.update()
+        print(device.info)
+
+        await led.upgrade(version="0.13.0-b4")
+
+
+if __name__ == "__main__":
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())

poetry.lock

@@ -31,6 +31,7 @@ aiohttp = ">=3.0.0"
 yarl = ">=1.6.0"
 backoff = ">=1.9.0"
 awesomeversion = ">=21.10.1"
+cachetools = ">=4.0.0"
 
 [tool.poetry.dev-dependencies]
 aresponses = "^2.1.4"
@@ -61,6 +62,7 @@ darglint = "^1.8.1"
 safety = "^1.10.3"
 codespell = "^2.1.0"
 bandit = "^1.7.0"
+types-cachetools = "^4.2.4"
 
 [tool.poetry.urls]
 "Bug Tracker" = "https://github.com/frenck/python-wled/issues"

pyproject.toml

@@ -5,6 +5,8 @@
 from enum import IntEnum
 from typing import Any
 
+from awesomeversion import AwesomeVersion
+
 from .exceptions import WLEDError
 
 
@@ -316,7 +318,9 @@ class Info:  # pylint: disable=too-many-instance-attributes
     udp_port: int
     uptime: int
     version_id: str
-    version: str
+    version: AwesomeVersion | None
+    version_latest_beta: AwesomeVersion | None
+    version_latest_stable: AwesomeVersion | None
     websocket: int | None
     wifi: Wifi | None
 
@@ -333,6 +337,15 @@ def from_dict(data: dict[str, Any]) -> Info:
         if (websocket := data.get("ws")) == -1:
             websocket = None
 
+        if version := data.get("ver"):
+            version = AwesomeVersion(version)
+
+        if version_latest_stable := data.get("version_latest_stable"):
+            version_latest_stable = AwesomeVersion(version_latest_stable)
+
+        if version_latest_beta := data.get("version_latest_beta"):
+            version_latest_beta = AwesomeVersion(version_latest_beta)
+
         return Info(
             architecture=data.get("arch", "Unknown"),
             arduino_core_version=data.get("core", "Unknown").replace("_", "."),
@@ -352,7 +365,9 @@ def from_dict(data: dict[str, Any]) -> Info:
             udp_port=data.get("udpport", 0),
             uptime=data.get("uptime", 0),
             version_id=data.get("vid", "Unknown"),
-            version=data.get("ver", "Unknown"),
+            version=version,
+            version_latest_beta=version_latest_beta,
+            version_latest_stable=version_latest_stable,
             websocket=websocket,
             wifi=Wifi.from_dict(data),
         )

src/wled/models.py

@@ -5,13 +5,15 @@
 import json
 import socket
 from collections.abc import Callable
+from contextlib import suppress
 from dataclasses import dataclass
 from typing import Any
 
 import aiohttp
 import async_timeout
 import backoff  # type: ignore
 from awesomeversion import AwesomeVersion, AwesomeVersionException
+from cachetools import TTLCache
 from yarl import URL
 
 from .exceptions import (
@@ -37,6 +39,7 @@ class WLED:
     _device: Device | None = None
     _supports_si_request: bool | None = None
     _supports_presets: bool | None = None
+    _version_cache: TTLCache = TTLCache(maxsize=16, ttl=7200)
 
     @property
     def connected(self) -> bool:
@@ -169,7 +172,7 @@ async def request(
             data["v"] = True
 
         try:
-            with async_timeout.timeout(self.request_timeout):
+            async with async_timeout.timeout(self.request_timeout):
                 response = await self.session.request(
                     method,
                     url,
@@ -240,14 +243,17 @@ async def update(self, full_update: bool = False) -> Device:
             except WLEDError:
                 self._supports_presets = False
 
+            versions = await self.get_wled_versions_from_github()
+            data["info"].update(versions)
+
             self._device = Device(data)
 
             # Try to figure out if this version supports
             # a single info and state call
             try:
-                current = AwesomeVersion(self._device.info.version)
-                supported = AwesomeVersion("0.10.0")
-                self._supports_si_request = current >= supported
+                self._supports_si_request = self._device.info.version >= AwesomeVersion(
+                    "0.10.0"
+                )
             except AwesomeVersionException:
                 # Could be a manual build one? Lets poll for it
                 try:
@@ -279,6 +285,10 @@ async def update(self, full_update: bool = False) -> Device:
                     f"WLED device {self.host} returned an empty API"
                     " response on state update"
                 )
+
+            versions = await self.get_wled_versions_from_github()
+            info.update(versions)
+
             self._device.update_from_dict({"info": info, "state": state})
             return self._device
 
@@ -287,6 +297,10 @@ async def update(self, full_update: bool = False) -> Device:
                 f"WLED device at {self.host} returned an empty API"
                 " response on state & info update"
             )
+
+        versions = await self.get_wled_versions_from_github()
+        state_info["info"].update(versions)
+
         self._device.update_from_dict(state_info)
 
         return self._device
@@ -569,6 +583,127 @@ async def nightlight(
 
         await self.request("/json/state", method="POST", data=state)
 
+    async def upgrade(self, *, version: str) -> None:
+        """Upgrades WLED device to the specified version.
+
+        Args:
+            version: The version to upgrade to.
+
+        Raises:
+            WLEDError: If the upgrade fails.
+            WLEDConnectionTimeoutError: When a connection timeout occurs.
+            WLEDConnectionError: When a connection error occurs.
+        """
+        if self._device is None:
+            await self.update()
+
+        if self.session is None or self._device is None:
+            return
+
+        if self._device.info.architecture not in {"esp8266", "esp32"}:
+            raise WLEDError("Upgrade is only supported on ESP8266 and ESP32")
+
+        url = URL.build(scheme="http", host=self.host, port=80, path="/update")
+        update_file = f"WLED_{version}_{self._device.info.architecture.upper()}.bin"
+        download_url = f"https://github.com/Aircoookie/WLED/releases/download/v{version}/{update_file}"
+
+        try:
+            async with async_timeout.timeout(self.request_timeout * 10):
+                async with self.session.get(
+                    download_url, raise_for_status=True
+                ) as download:
+                    form = aiohttp.FormData()
+                    form.add_field("file", await download.read(), filename=update_file)
+                    await self.session.post(url, data=form)
+        except asyncio.TimeoutError as exception:
+            raise WLEDConnectionTimeoutError(
+                "Timeout occurred while fetching WLED version information from GitHub"
+            ) from exception
+        except aiohttp.ClientResponseError as exception:
+            if exception.status == 404:
+                raise WLEDError(
+                    f"Requested WLED version '{version}' does not exists"
+                ) from exception
+            raise WLEDError(
+                f"Could not download requested WLED version '{version}' from {download_url}"
+            ) from exception
+        except (aiohttp.ClientError, socket.gaierror) as exception:
+            raise WLEDConnectionError(
+                "Timeout occurred while communicating with GitHub for WLED version information"
+            ) from exception
+
+    @backoff.on_exception(backoff.expo, WLEDConnectionError, max_tries=3, logger=None)
+    async def get_wled_versions_from_github(self) -> dict[str, str | None]:
+        """Fetch WLED version information from GitHub.
+
+        Returns:
+            A dictionary of WLED versions, with the key being the version type.
+
+        Raises:
+            WLEDConnectionTimeoutError: Timeout occurred while fetching WLED
+                version information from GitHub.
+            WLEDConnectionError: Timeout occurred while communicating with
+                GitHub for WLED version information.
+            WLEDError: Didn't get a JSON response from GitHub while retrieving
+                version information.
+        """
+        with suppress(KeyError):
+            return {
+                "version_latest_stable": self._version_cache["stable"],
+                "version_latest_beta": self._version_cache["beta"],
+            }
+
+        if self.session is None:
+            return {"version_latest_stable": None, "version_latest_beta": None}
+
+        try:
+            async with async_timeout.timeout(self.request_timeout):
+                response = await self.session.get(
+                    "https://api.github.com/repos/Aircoookie/WLED/releases"
+                )
+        except asyncio.TimeoutError as exception:
+            raise WLEDConnectionTimeoutError(
+                "Timeout occurred while fetching WLED version information from GitHub"
+            ) from exception
+        except (aiohttp.ClientError, socket.gaierror) as exception:
+            raise WLEDConnectionError(
+                "Timeout occurred while communicating with GitHub for WLED version"
+            ) from exception
+
+        content_type = response.headers.get("Content-Type", "")
+        if (response.status // 100) in [4, 5]:
+            contents = await response.read()
+            response.close()
+
+            if content_type == "application/json":
+                raise WLEDError(response.status, json.loads(contents.decode("utf8")))
+            raise WLEDError(response.status, {"message": contents.decode("utf8")})
+
+        if "application/json" not in content_type:
+            raise WLEDError(
+                "Didn't get a JSON response from GitHub while retrieving version information"
+            )
+
+        releases = await response.json()
+        version_latest = None
+        version_latest_beta = None
+        for release in releases:
+            if release["prerelease"] is False and version_latest is None:
+                version_latest = release["tag_name"].lstrip("vV")
+            if release["prerelease"] is True and version_latest_beta is None:
+                version_latest_beta = release["tag_name"].lstrip("vV")
+            if version_latest is not None and version_latest_beta is not None:
+                break
+
+        # Cache results
+        self._version_cache["stable"] = version_latest
+        self._version_cache["beta"] = version_latest_beta
+
+        return {
+            "version_latest_stable": version_latest,
+            "version_latest_beta": version_latest_beta,
+        }
+
     async def reset(self) -> None:
         """Reboot WLED device."""
         await self.request("/reset")

src/wled/wled.py

@@ -1,5 +1,6 @@
 """Tests for `wled.WLED`."""
 import asyncio
+from unittest.mock import patch
 
 import aiohttp
 import pytest
@@ -8,6 +9,16 @@
 from wled.exceptions import WLEDConnectionError, WLEDError
 
 
+@pytest.fixture(autouse=True)
+def mock_get_version_from_github():
+    """Patch out connection to GitHub."""
+    with patch(
+        "wled.WLED.get_wled_versions_from_github",
+        return_value={"version_latest_stable": None, "version_latest_beta": None},
+    ):
+        yield
+
+
 @pytest.mark.asyncio
 async def test_json_request(aresponses):
     """Test JSON response is handled correctly."""