Add update platform to WLED (#68454)

* Add update platform to WLED

* Copy pasta fixes

* Apply suggestions from code review

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update tests/components/wled/test_update.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Fix tests

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/wled/__init__.py

@@ -16,6 +16,7 @@ PLATFORMS = (
     Platform.SELECT,
     Platform.SENSOR,
     Platform.SWITCH,
+    Platform.UPDATE,
 )
 
 

homeassistant/components/wled/binary_sensor.py

@@ -35,6 +35,9 @@ class WLEDUpdateBinarySensor(WLEDEntity, BinarySensorEntity):
     _attr_entity_category = EntityCategory.DIAGNOSTIC
     _attr_device_class = BinarySensorDeviceClass.UPDATE
 
+    # Disabled by default, as this entity is deprecated.
+    _attr_entity_registry_enabled_default = False
+
     def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize the button entity."""
         super().__init__(coordinator=coordinator)

homeassistant/components/wled/button.py

@@ -7,7 +7,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from .const import DOMAIN, LOGGER
 from .coordinator import WLEDDataUpdateCoordinator
 from .helpers import wled_exception_handler
 from .models import WLEDEntity
@@ -52,6 +52,9 @@ class WLEDUpdateButton(WLEDEntity, ButtonEntity):
     _attr_device_class = ButtonDeviceClass.UPDATE
     _attr_entity_category = EntityCategory.CONFIG
 
+    # Disabled by default, as this entity is deprecated.
+    _attr_entity_registry_enabled_default = False
+
     def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize the button entity."""
         super().__init__(coordinator=coordinator)
@@ -83,6 +86,11 @@ class WLEDUpdateButton(WLEDEntity, ButtonEntity):
     @wled_exception_handler
     async def async_press(self) -> None:
         """Send out a update command."""
+        LOGGER.warning(
+            "The WLED update button '%s' is deprecated, please "
+            "use the new update entity as a replacement",
+            self.entity_id,
+        )
         current = self.coordinator.data.info.version
         beta = self.coordinator.data.info.version_latest_beta
         stable = self.coordinator.data.info.version_latest_stable

homeassistant/components/wled/update.py

@@ -0,0 +1,93 @@
+"""Support for WLED updates."""
+from __future__ import annotations
+
+from typing import Any, cast
+
+from homeassistant.components.update import (
+    UpdateDeviceClass,
+    UpdateEntity,
+    UpdateEntityFeature,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import DOMAIN
+from .coordinator import WLEDDataUpdateCoordinator
+from .helpers import wled_exception_handler
+from .models import WLEDEntity
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up WLED update based on a config entry."""
+    coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities([WLEDUpdateEntity(coordinator)])
+
+
+class WLEDUpdateEntity(WLEDEntity, UpdateEntity):
+    """Defines a WLED update entity."""
+
+    _attr_device_class = UpdateDeviceClass.FIRMWARE
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL | UpdateEntityFeature.SPECIFIC_VERSION
+    )
+
+    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
+        """Initialize the update entity."""
+        super().__init__(coordinator=coordinator)
+        self._attr_name = f"{coordinator.data.info.name} Firmware"
+        self._attr_unique_id = coordinator.data.info.mac_address
+        self._attr_title = "WLED"
+
+    @property
+    def current_version(self) -> str | None:
+        """Version currently in use."""
+        if (version := self.coordinator.data.info.version) is None:
+            return None
+        return str(version)
+
+    @property
+    def latest_version(self) -> str | None:
+        """Latest version available for install."""
+        # If we already run a pre-release, we consider being on the beta channel.
+        # Offer beta version upgrade, unless stable is newer
+        if (
+            (beta := self.coordinator.data.info.version_latest_beta) is not None
+            and (current := self.coordinator.data.info.version) is not None
+            and (current.alpha or current.beta or current.release_candidate)
+            and (
+                (stable := self.coordinator.data.info.version_latest_stable) is None
+                or (stable is not None and stable < beta)
+            )
+        ):
+            return str(beta)
+
+        if (stable := self.coordinator.data.info.version_latest_stable) is not None:
+            return str(stable)
+
+        return None
+
+    @property
+    def release_url(self) -> str | None:
+        """URL to the full release notes of the latest version available."""
+        if (version := self.latest_version) is None:
+            return None
+        return f"https://github.com/Aircoookie/WLED/releases/tag/v{version}"
+
+    @wled_exception_handler
+    async def async_install(
+        self,
+        version: str | None = None,
+        backup: bool | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Install an update."""
+        if version is None:
+            # We cast here, as we know that the latest_version is a string.
+            version = cast(str, self.latest_version)
+        await self.coordinator.wled.upgrade(version=version)
+        await self.coordinator.async_refresh()

tests/components/wled/fixtures/rgb_no_update.json

@@ -0,0 +1,218 @@
+{
+  "state": {
+    "on": true,
+    "bri": 127,
+    "transition": 7,
+    "ps": -1,
+    "pl": -1,
+    "nl": {
+      "on": false,
+      "dur": 60,
+      "fade": true,
+      "tbri": 0
+    },
+    "udpn": {
+      "send": false,
+      "recv": true
+    },
+    "seg": [
+      {
+        "id": 0,
+        "start": 0,
+        "stop": 19,
+        "len": 20,
+        "col": [[255, 159, 0], [0, 0, 0], [0, 0, 0]],
+        "fx": 0,
+        "sx": 32,
+        "ix": 128,
+        "pal": 0,
+        "sel": true,
+        "rev": false,
+        "cln": -1
+      },
+      {
+        "id": 1,
+        "start": 20,
+        "stop": 30,
+        "len": 10,
+        "col": [[0, 255, 123], [0, 0, 0], [0, 0, 0]],
+        "fx": 1,
+        "sx": 16,
+        "ix": 64,
+        "pal": 1,
+        "sel": true,
+        "rev": true,
+        "cln": -1
+      }
+    ]
+  },
+  "info": {
+    "ver": null,
+    "version_latest_stable": null,
+    "version_latest_beta": null,
+    "vid": 1909122,
+    "leds": {
+      "count": 30,
+      "rgbw": false,
+      "pin": [2],
+      "pwr": 470,
+      "maxpwr": 850,
+      "maxseg": 10
+    },
+    "name": "WLED RGB Light",
+    "udpport": 21324,
+    "live": false,
+    "fxcount": 81,
+    "palcount": 50,
+    "wifi": {
+      "bssid": "AA:AA:AA:AA:AA:BB",
+      "rssi": -62,
+      "signal": 76,
+      "channel": 11
+    },
+    "arch": "esp8266",
+    "core": "2_4_2",
+    "freeheap": 14600,
+    "uptime": 32,
+    "opt": 119,
+    "brand": "WLED",
+    "product": "DIY light",
+    "btype": "bin",
+    "mac": "aabbccddeeff"
+  },
+  "effects": [
+    "Solid",
+    "Blink",
+    "Breathe",
+    "Wipe",
+    "Wipe Random",
+    "Random Colors",
+    "Sweep",
+    "Dynamic",
+    "Colorloop",
+    "Rainbow",
+    "Scan",
+    "Dual Scan",
+    "Fade",
+    "Chase",
+    "Chase Rainbow",
+    "Running",
+    "Saw",
+    "Twinkle",
+    "Dissolve",
+    "Dissolve Rnd",
+    "Sparkle",
+    "Dark Sparkle",
+    "Sparkle+",
+    "Strobe",
+    "Strobe Rainbow",
+    "Mega Strobe",
+    "Blink Rainbow",
+    "Android",
+    "Chase",
+    "Chase Random",
+    "Chase Rainbow",
+    "Chase Flash",
+    "Chase Flash Rnd",
+    "Rainbow Runner",
+    "Colorful",
+    "Traffic Light",
+    "Sweep Random",
+    "Running 2",
+    "Red & Blue",
+    "Stream",
+    "Scanner",
+    "Lighthouse",
+    "Fireworks",
+    "Rain",
+    "Merry Christmas",
+    "Fire Flicker",
+    "Gradient",
+    "Loading",
+    "In Out",
+    "In In",
+    "Out Out",
+    "Out In",
+    "Circus",
+    "Halloween",
+    "Tri Chase",
+    "Tri Wipe",
+    "Tri Fade",
+    "Lightning",
+    "ICU",
+    "Multi Comet",
+    "Dual Scanner",
+    "Stream 2",
+    "Oscillate",
+    "Pride 2015",
+    "Juggle",
+    "Palette",
+    "Fire 2012",
+    "Colorwaves",
+    "BPM",
+    "Fill Noise",
+    "Noise 1",
+    "Noise 2",
+    "Noise 3",
+    "Noise 4",
+    "Colortwinkle",
+    "Lake",
+    "Meteor",
+    "Smooth Meteor",
+    "Railway",
+    "Ripple",
+    "Twinklefox"
+  ],
+  "palettes": [
+    "Default",
+    "Random Cycle",
+    "Primary Color",
+    "Based on Primary",
+    "Set Colors",
+    "Based on Set",
+    "Party",
+    "Cloud",
+    "Lava",
+    "Ocean",
+    "Forest",
+    "Rainbow",
+    "Rainbow Bands",
+    "Sunset",
+    "Rivendell",
+    "Breeze",
+    "Red & Blue",
+    "Yellowout",
+    "Analogous",
+    "Splash",
+    "Pastel",
+    "Sunset 2",
+    "Beech",
+    "Vintage",
+    "Departure",
+    "Landscape",
+    "Beach",
+    "Sherbet",
+    "Hult",
+    "Hult 64",
+    "Drywet",
+    "Jul",
+    "Grintage",
+    "Rewhi",
+    "Tertiary",
+    "Fire",
+    "Icefire",
+    "Cyane",
+    "Light Pink",
+    "Autumn",
+    "Magenta",
+    "Magred",
+    "Yelmag",
+    "Yelblu",
+    "Orange & Teal",
+    "Tiamat",
+    "April Night",
+    "Orangery",
+    "C9",
+    "Sakura"
+  ]
+}

tests/components/wled/test_binary_sensor.py

@@ -1,5 +1,5 @@
 """Tests for the WLED binary sensor platform."""
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -13,7 +13,10 @@ from tests.common import MockConfigEntry
 
 
 async def test_update_available(
-    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+    hass: HomeAssistant,
+    entity_registry_enabled_by_default: AsyncMock,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
 ) -> None:
     """Test the firmware update binary sensor."""
     entity_registry = er.async_get(hass)
@@ -32,7 +35,10 @@ async def test_update_available(
 
 @pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True)
 async def test_no_update_available(
-    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+    hass: HomeAssistant,
+    entity_registry_enabled_by_default: AsyncMock,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
 ) -> None:
     """Test the update binary sensor. There is no update available."""
     entity_registry = er.async_get(hass)
@@ -47,3 +53,18 @@ async def test_no_update_available(
     assert entry
     assert entry.unique_id == "aabbccddeeff_update"
     assert entry.entity_category is EntityCategory.DIAGNOSTIC
+
+
+async def test_disabled_by_default(
+    hass: HomeAssistant, init_integration: MockConfigEntry
+) -> None:
+    """Test that the binary update sensor is disabled by default."""
+    registry = er.async_get(hass)
+
+    state = hass.states.get("binary_sensor.wled_rgb_light_firmware")
+    assert state is None
+
+    entry = registry.async_get("binary_sensor.wled_rgb_light_firmware")
+    assert entry
+    assert entry.disabled
+    assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION

tests/components/wled/test_button.py

@@ -1,5 +1,5 @@
 """Tests for the WLED button platform."""
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 from freezegun import freeze_time
 import pytest
@@ -98,7 +98,11 @@ async def test_button_connection_error(
 
 
 async def test_button_update_stay_stable(
-    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+    hass: HomeAssistant,
+    entity_registry_enabled_by_default: AsyncMock,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+    caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test the update button.
 
@@ -127,11 +131,19 @@ async def test_button_update_stay_stable(
     await hass.async_block_till_done()
     assert mock_wled.upgrade.call_count == 1
     mock_wled.upgrade.assert_called_with(version="0.12.0")
+    assert (
+        "The WLED update button 'button.wled_rgb_light_update' is deprecated"
+        in caplog.text
+    )
 
 
 @pytest.mark.parametrize("mock_wled", ["wled/rgbw.json"], indirect=True)
 async def test_button_update_beta_to_stable(
-    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+    hass: HomeAssistant,
+    entity_registry_enabled_by_default: AsyncMock,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+    caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test the update button.
 
@@ -148,11 +160,19 @@ async def test_button_update_beta_to_stable(
     await hass.async_block_till_done()
     assert mock_wled.upgrade.call_count == 1
     mock_wled.upgrade.assert_called_with(version="0.8.6")
+    assert (
+        "The WLED update button 'button.wled_rgbw_light_update' is deprecated"
+        in caplog.text
+    )
 
 
 @pytest.mark.parametrize("mock_wled", ["wled/rgb_single_segment.json"], indirect=True)
 async def test_button_update_stay_beta(
-    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+    hass: HomeAssistant,
+    entity_registry_enabled_by_default: AsyncMock,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+    caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test the update button.
 
@@ -168,13 +188,35 @@ async def test_button_update_stay_beta(
     await hass.async_block_till_done()
     assert mock_wled.upgrade.call_count == 1
     mock_wled.upgrade.assert_called_with(version="0.8.6b2")
+    assert (
+        "The WLED update button 'button.wled_rgb_light_update' is deprecated"
+        in caplog.text
+    )
 
 
 @pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True)
 async def test_button_no_update_available(
-    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+    hass: HomeAssistant,
+    entity_registry_enabled_by_default: AsyncMock,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
 ) -> None:
     """Test the update button. There is no update available."""
     state = hass.states.get("button.wled_websocket_update")
     assert state
     assert state.state == STATE_UNAVAILABLE
+
+
+async def test_disabled_by_default(
+    hass: HomeAssistant, init_integration: MockConfigEntry
+) -> None:
+    """Test that the update button is disabled by default."""
+    registry = er.async_get(hass)
+
+    state = hass.states.get("button.wled_rgb_light_update")
+    assert state is None
+
+    entry = registry.async_get("button.wled_rgb_light_update")
+    assert entry
+    assert entry.disabled
+    assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION

tests/components/wled/test_update.py

@@ -0,0 +1,240 @@
+"""Tests for the WLED update platform."""
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from wled import WLEDError
+
+from homeassistant.components.update import (
+    DOMAIN as UPDATE_DOMAIN,
+    SERVICE_INSTALL,
+    UpdateDeviceClass,
+    UpdateEntityFeature,
+)
+from homeassistant.components.update.const import (
+    ATTR_CURRENT_VERSION,
+    ATTR_LATEST_VERSION,
+    ATTR_RELEASE_SUMMARY,
+    ATTR_RELEASE_URL,
+    ATTR_TITLE,
+)
+from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
+    ATTR_ENTITY_ID,
+    ATTR_ICON,
+    ATTR_SUPPORTED_FEATURES,
+    STATE_OFF,
+    STATE_ON,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.entity import EntityCategory
+
+from tests.common import MockConfigEntry
+
+
+async def test_update_available(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+) -> None:
+    """Test the firmware update available."""
+    entity_registry = er.async_get(hass)
+
+    state = hass.states.get("update.wled_rgb_light_firmware")
+    assert state
+    assert state.attributes.get(ATTR_DEVICE_CLASS) == UpdateDeviceClass.FIRMWARE
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_CURRENT_VERSION] == "0.8.5"
+    assert state.attributes[ATTR_LATEST_VERSION] == "0.12.0"
+    assert state.attributes[ATTR_RELEASE_SUMMARY] is None
+    assert (
+        state.attributes[ATTR_RELEASE_URL]
+        == "https://github.com/Aircoookie/WLED/releases/tag/v0.12.0"
+    )
+    assert (
+        state.attributes[ATTR_SUPPORTED_FEATURES]
+        == UpdateEntityFeature.INSTALL | UpdateEntityFeature.SPECIFIC_VERSION
+    )
+    assert state.attributes[ATTR_TITLE] == "WLED"
+    assert ATTR_ICON not in state.attributes
+
+    entry = entity_registry.async_get("update.wled_rgb_light_firmware")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff"
+    assert entry.entity_category is EntityCategory.CONFIG
+
+
+@pytest.mark.parametrize("mock_wled", ["wled/rgb_no_update.json"], indirect=True)
+async def test_update_information_available(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+) -> None:
+    """Test having no update information available at all."""
+    entity_registry = er.async_get(hass)
+
+    state = hass.states.get("update.wled_rgb_light_firmware")
+    assert state
+    assert state.attributes.get(ATTR_DEVICE_CLASS) == UpdateDeviceClass.FIRMWARE
+    assert state.state == STATE_UNKNOWN
+    assert state.attributes[ATTR_CURRENT_VERSION] is None
+    assert state.attributes[ATTR_LATEST_VERSION] is None
+    assert state.attributes[ATTR_RELEASE_SUMMARY] is None
+    assert state.attributes[ATTR_RELEASE_URL] is None
+    assert (
+        state.attributes[ATTR_SUPPORTED_FEATURES]
+        == UpdateEntityFeature.INSTALL | UpdateEntityFeature.SPECIFIC_VERSION
+    )
+    assert state.attributes[ATTR_TITLE] == "WLED"
+    assert ATTR_ICON not in state.attributes
+
+    entry = entity_registry.async_get("update.wled_rgb_light_firmware")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff"
+    assert entry.entity_category is EntityCategory.CONFIG
+
+
+@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True)
+async def test_no_update_available(
+    hass: HomeAssistant,
+    entity_registry_enabled_by_default: AsyncMock,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+) -> None:
+    """Test there is no update available."""
+    entity_registry = er.async_get(hass)
+
+    state = hass.states.get("update.wled_websocket_firmware")
+    assert state
+    assert state.state == STATE_OFF
+    assert state.attributes.get(ATTR_DEVICE_CLASS) == UpdateDeviceClass.FIRMWARE
+    assert state.attributes[ATTR_CURRENT_VERSION] == "0.12.0-b2"
+    assert state.attributes[ATTR_LATEST_VERSION] == "0.12.0-b2"
+    assert state.attributes[ATTR_RELEASE_SUMMARY] is None
+    assert (
+        state.attributes[ATTR_RELEASE_URL]
+        == "https://github.com/Aircoookie/WLED/releases/tag/v0.12.0-b2"
+    )
+    assert (
+        state.attributes[ATTR_SUPPORTED_FEATURES]
+        == UpdateEntityFeature.INSTALL | UpdateEntityFeature.SPECIFIC_VERSION
+    )
+    assert state.attributes[ATTR_TITLE] == "WLED"
+    assert ATTR_ICON not in state.attributes
+
+    assert ATTR_ICON not in state.attributes
+
+    entry = entity_registry.async_get("update.wled_websocket_firmware")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff"
+    assert entry.entity_category is EntityCategory.CONFIG
+
+
+async def test_update_error(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test error handling of the WLED update."""
+    mock_wled.update.side_effect = WLEDError
+
+    await hass.services.async_call(
+        UPDATE_DOMAIN,
+        SERVICE_INSTALL,
+        {ATTR_ENTITY_ID: "update.wled_rgb_light_firmware"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("update.wled_rgb_light_firmware")
+    assert state
+    assert state.state == STATE_UNAVAILABLE
+    assert "Invalid response from API" in caplog.text
+
+
+async def test_update_stay_stable(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+) -> None:
+    """Test the update entity staying on stable.
+
+    There is both an update for beta and stable available, however, the device
+    is currently running a stable version. Therefore, the update entity should
+    update to the next stable (even though beta is newer).
+    """
+    state = hass.states.get("update.wled_rgb_light_firmware")
+    assert state
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_CURRENT_VERSION] == "0.8.5"
+    assert state.attributes[ATTR_LATEST_VERSION] == "0.12.0"
+
+    await hass.services.async_call(
+        UPDATE_DOMAIN,
+        SERVICE_INSTALL,
+        {ATTR_ENTITY_ID: "update.wled_rgb_light_firmware"},
+        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_update_beta_to_stable(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+) -> None:
+    """Test the update entity.
+
+    There is both an update for beta and stable available and the device
+    is currently a beta, however, a newer stable is available. Therefore, the
+    update entity should update to the next stable.
+    """
+    state = hass.states.get("update.wled_rgbw_light_firmware")
+    assert state
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_CURRENT_VERSION] == "0.8.6b4"
+    assert state.attributes[ATTR_LATEST_VERSION] == "0.8.6"
+
+    await hass.services.async_call(
+        UPDATE_DOMAIN,
+        SERVICE_INSTALL,
+        {ATTR_ENTITY_ID: "update.wled_rgbw_light_firmware"},
+        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_update_stay_beta(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+) -> None:
+    """Test the update entity.
+
+    There is an update for beta and the device is currently a beta. Therefore,
+    the update entity should update to the next beta.
+    """
+    state = hass.states.get("update.wled_rgb_light_firmware")
+    assert state
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_CURRENT_VERSION] == "0.8.6b1"
+    assert state.attributes[ATTR_LATEST_VERSION] == "0.8.6b2"
+
+    await hass.services.async_call(
+        UPDATE_DOMAIN,
+        SERVICE_INSTALL,
+        {ATTR_ENTITY_ID: "update.wled_rgb_light_firmware"},
+        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")