Add WLED firmware upgrade button (#59793)

homeassistant/components/wled/button.py

@@ -20,11 +20,16 @@ async def async_setup_entry(
 ) -> None:
     """Set up WLED button based on a config entry."""
     coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([WLEDRestartButton(coordinator)])
+    async_add_entities(
+        [
+            WLEDRestartButton(coordinator),
+            WLEDUpgradeButton(coordinator),
+        ]
+    )
 
 
 class WLEDRestartButton(WLEDEntity, ButtonEntity):
-    """Defines a WLED restart switch."""
+    """Defines a WLED restart button."""
 
     _attr_icon = "mdi:restart"
     _attr_entity_category = ENTITY_CATEGORY_CONFIG
@@ -39,3 +44,59 @@ class WLEDRestartButton(WLEDEntity, ButtonEntity):
     async def async_press(self) -> None:
         """Send out a restart command."""
         await self.coordinator.wled.reset()
+
+
+class WLEDUpgradeButton(WLEDEntity, ButtonEntity):
+    """Defines a WLED upgrade button."""
+
+    _attr_icon = "mdi:cellphone-arrow-down"
+    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+
+    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
+        """Initialize the button entity."""
+        super().__init__(coordinator=coordinator)
+        self._attr_name = f"{coordinator.data.info.name} Upgrade"
+        self._attr_unique_id = f"{coordinator.data.info.mac_address}_upgrade"
+
+    @property
+    def available(self) -> bool:
+        """Return if the entity and an upgrade is available."""
+        current = self.coordinator.data.info.version
+        beta = self.coordinator.data.info.version_latest_beta
+        stable = self.coordinator.data.info.version_latest_stable
+
+        # If we already run a pre-release, allow upgrading to a newer
+        # pre-release offer a normal upgrade otherwise.
+        return (
+            super().available
+            and current is not None
+            and (
+                (stable is not None and stable > current)
+                or (
+                    beta is not None
+                    and (current.alpha or current.beta or current.release_candidate)
+                    and beta > current
+                )
+            )
+        )
+
+    @wled_exception_handler
+    async def async_press(self) -> None:
+        """Send out a restart command."""
+        current = self.coordinator.data.info.version
+        beta = self.coordinator.data.info.version_latest_beta
+        stable = self.coordinator.data.info.version_latest_stable
+
+        # If we already run a pre-release, allow upgrading to a newer
+        # pre-release or newer stable, otherwise, offer a normal stable upgrades.
+        version = stable
+        if (
+            current is not None
+            and beta is not None
+            and (current.alpha or current.beta or current.release_candidate)
+            and beta > current
+            and beta > stable
+        ):
+            version = beta
+
+        await self.coordinator.wled.upgrade(version=str(version))