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))

tests/components/wled/fixtures/rgb.json

@@ -48,6 +48,8 @@
   },
   "info": {
     "ver": "0.8.5",
+    "version_latest_stable": "0.12.0",
+    "version_latest_beta": "0.13.0b1",
     "vid": 1909122,
     "leds": {
       "count": 30,

tests/components/wled/fixtures/rgb_single_segment.json

@@ -33,7 +33,9 @@
     ]
   },
   "info": {
-    "ver": "0.8.5",
+    "ver": "0.8.6b1",
+    "version_latest_stable": "0.8.5",
+    "version_latest_beta": "0.8.6b2",
     "vid": 1909122,
     "leds": {
       "count": 30,

tests/components/wled/fixtures/rgb_websocket.json

@@ -63,6 +63,8 @@
   },
   "info": {
       "ver": "0.12.0-b2",
+      "version_latest_stable": "0.11.0",
+      "version_latest_beta": "0.12.0-b2",
       "vid": 2103220,
       "leds": {
           "count": 13,

tests/components/wled/fixtures/rgbw.json

@@ -33,7 +33,9 @@
     ]
   },
   "info": {
-    "ver": "0.8.6",
+    "ver": "0.8.6b4",
+    "version_latest_stable": "0.8.6",
+    "version_latest_beta": "0.8.6b5",
     "vid": 1910255,
     "leds": {
       "count": 13,

tests/components/wled/test_button.py

@@ -19,7 +19,7 @@ from homeassistant.helpers import entity_registry as er
 from tests.common import MockConfigEntry
 
 
-async def test_button(
+async def test_button_restart(
     hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
 ) -> None:
     """Test the creation and values of the WLED button."""
@@ -35,7 +35,6 @@ async def test_button(
     assert entry.unique_id == "aabbccddeeff_restart"
     assert entry.entity_category == ENTITY_CATEGORY_CONFIG
 
-    # Restart
     await hass.services.async_call(
         BUTTON_DOMAIN,
         SERVICE_PRESS,
@@ -92,3 +91,86 @@ async def test_button_connection_error(
     assert state
     assert state.state == STATE_UNAVAILABLE
     assert "Error communicating with API" in caplog.text
+
+
+async def test_button_upgrade_stay_stable(
+    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+) -> None:
+    """Test the upgrade button.
+
+    There is both an upgrade for beta and stable available, however, the device
+    is currently running a stable version. Therefore, the upgrade button should
+    upgrade the the next stable (even though beta is newer).
+    """
+    entity_registry = er.async_get(hass)
+
+    entry = entity_registry.async_get("button.wled_rgb_light_upgrade")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_upgrade"
+    assert entry.entity_category == ENTITY_CATEGORY_CONFIG
+
+    state = hass.states.get("button.wled_rgb_light_upgrade")
+    assert state
+    assert state.attributes.get(ATTR_ICON) == "mdi:cellphone-arrow-down"
+    assert state.state == STATE_UNKNOWN
+
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.wled_rgb_light_upgrade"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert mock_wled.upgrade.call_count == 1
+    mock_wled.upgrade.assert_called_with(version="0.12.0")
+
+
+@pytest.mark.parametrize("mock_wled", ["wled/rgbw.json"], indirect=True)
+async def test_button_upgrade_beta_to_stable(
+    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+) -> None:
+    """Test the upgrade button.
+
+    There is both an upgrade for beta and stable available the device
+    is currently a beta, however, a newer stable is available. Therefore, the
+    upgrade button should upgrade to the next stable.
+    """
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.wled_rgbw_light_upgrade"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert mock_wled.upgrade.call_count == 1
+    mock_wled.upgrade.assert_called_with(version="0.8.6")
+
+
+@pytest.mark.parametrize("mock_wled", ["wled/rgb_single_segment.json"], indirect=True)
+async def test_button_upgrade_stay_beta(
+    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+) -> None:
+    """Test the upgrade button.
+
+    There is an upgrade for beta and the device is currently a beta. Therefore,
+    the upgrade button should upgrade to the next beta.
+    """
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.wled_rgb_light_upgrade"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert mock_wled.upgrade.call_count == 1
+    mock_wled.upgrade.assert_called_with(version="0.8.6b2")
+
+
+@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True)
+async def test_button_no_upgrade_available(
+    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+) -> None:
+    """Test the upgrade button. There is no update available."""
+    state = hass.states.get("button.wled_websocket_upgrade")
+    assert state
+    assert state.state == STATE_UNAVAILABLE