Shelly migrate to update entity (#78305)

* Add update entity

* fixes

* fixes

* change to CONFIG catogory

* return latest version if no update available

* fixes

* Remove firmware binary_sensors and buttons

* import Callable from collections

* remove ota_update tests

* Update homeassistant/components/shelly/update.py

Co-authored-by: Shay Levy <levyshay1@gmail.com>

* simplify

* fix mypy

* Create test_update.py

* fix isort

* add progress support

* fix styling

* fix update_tests

* fix styling

* do not exclude shelly update test

* bring coverage to 100%

* snake case

* snake case

* change str(x) to cast(str, x)

* simplify tests

* further simplify tests

* Split MOCK_SHELLY_COAP and MOCK_SHELLY_RPC

* fix issort

* fix status test

* fix isort

* run python3 -m script.hassfest

Co-authored-by: Shay Levy <levyshay1@gmail.com>

homeassistant/components/shelly/__init__.py

@@ -80,6 +80,7 @@ BLOCK_PLATFORMS: Final = [
     Platform.LIGHT,
     Platform.SENSOR,
     Platform.SWITCH,
+    Platform.UPDATE,
 ]
 BLOCK_SLEEPING_PLATFORMS: Final = [
     Platform.BINARY_SENSOR,
@@ -94,6 +95,7 @@ RPC_PLATFORMS: Final = [
     Platform.LIGHT,
     Platform.SENSOR,
     Platform.SWITCH,
+    Platform.UPDATE,
 ]
 
 

homeassistant/components/shelly/binary_sensor.py

@@ -143,19 +143,6 @@ REST_SENSORS: Final = {
         entity_registry_enabled_default=False,
         entity_category=EntityCategory.DIAGNOSTIC,
     ),
-    "fwupdate": RestBinarySensorDescription(
-        key="fwupdate",
-        name="Firmware Update",
-        device_class=BinarySensorDeviceClass.UPDATE,
-        value=lambda status, _: status["update"]["has_update"],
-        entity_registry_enabled_default=False,
-        extra_state_attributes=lambda status: {
-            "latest_stable_version": status["update"]["new_version"],
-            "installed_version": status["update"]["old_version"],
-            "beta_version": status["update"].get("beta_version", ""),
-        },
-        entity_category=EntityCategory.DIAGNOSTIC,
-    ),
 }
 
 RPC_SENSORS: Final = {
@@ -175,19 +162,6 @@ RPC_SENSORS: Final = {
         entity_registry_enabled_default=False,
         entity_category=EntityCategory.DIAGNOSTIC,
     ),
-    "fwupdate": RpcBinarySensorDescription(
-        key="sys",
-        sub_key="available_updates",
-        name="Firmware Update",
-        device_class=BinarySensorDeviceClass.UPDATE,
-        entity_registry_enabled_default=False,
-        extra_state_attributes=lambda status, shelly: {
-            "latest_stable_version": status.get("stable", {"version": ""})["version"],
-            "installed_version": shelly["ver"],
-            "beta_version": status.get("beta", {"version": ""})["version"],
-        },
-        entity_category=EntityCategory.DIAGNOSTIC,
-    ),
     "overtemp": RpcBinarySensorDescription(
         key="switch",
         sub_key="errors",

homeassistant/components/shelly/button.py

@@ -37,21 +37,6 @@ class ShellyButtonDescription(ButtonEntityDescription, ShellyButtonDescriptionMi
 
 
 BUTTONS: Final = [
-    ShellyButtonDescription(
-        key="ota_update",
-        name="OTA Update",
-        device_class=ButtonDeviceClass.UPDATE,
-        entity_category=EntityCategory.CONFIG,
-        press_action=lambda wrapper: wrapper.async_trigger_ota_update(),
-    ),
-    ShellyButtonDescription(
-        key="ota_update_beta",
-        name="OTA Update Beta",
-        device_class=ButtonDeviceClass.UPDATE,
-        entity_registry_enabled_default=False,
-        entity_category=EntityCategory.CONFIG,
-        press_action=lambda wrapper: wrapper.async_trigger_ota_update(beta=True),
-    ),
     ShellyButtonDescription(
         key="reboot",
         name="Reboot",

homeassistant/components/shelly/update.py

@@ -0,0 +1,236 @@
+"""Update entities for Shelly devices."""
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+import logging
+from typing import Any, Final, cast
+
+from homeassistant.components.update import (
+    UpdateDeviceClass,
+    UpdateEntity,
+    UpdateEntityDescription,
+    UpdateEntityFeature,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from . import BlockDeviceWrapper, RpcDeviceWrapper
+from .const import BLOCK, CONF_SLEEP_PERIOD, DATA_CONFIG_ENTRY, DOMAIN
+from .entity import (
+    RestEntityDescription,
+    RpcEntityDescription,
+    ShellyRestAttributeEntity,
+    ShellyRpcAttributeEntity,
+    async_setup_entry_rest,
+    async_setup_entry_rpc,
+)
+from .utils import get_device_entry_gen
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class RpcUpdateRequiredKeysMixin:
+    """Class for RPC update required keys."""
+
+    latest_version: Callable[[dict], Any]
+    install: Callable
+
+
+@dataclass
+class RestUpdateRequiredKeysMixin:
+    """Class for REST update required keys."""
+
+    latest_version: Callable[[dict], Any]
+    install: Callable
+
+
+@dataclass
+class RpcUpdateDescription(
+    RpcEntityDescription, UpdateEntityDescription, RpcUpdateRequiredKeysMixin
+):
+    """Class to describe a RPC update."""
+
+
+@dataclass
+class RestUpdateDescription(
+    RestEntityDescription, UpdateEntityDescription, RestUpdateRequiredKeysMixin
+):
+    """Class to describe a REST update."""
+
+
+REST_UPDATES: Final = {
+    "fwupdate": RestUpdateDescription(
+        name="Firmware Update",
+        key="fwupdate",
+        latest_version=lambda status: status["update"]["new_version"],
+        install=lambda wrapper: wrapper.async_trigger_ota_update(),
+        device_class=UpdateDeviceClass.FIRMWARE,
+        entity_category=EntityCategory.CONFIG,
+        entity_registry_enabled_default=True,
+    ),
+    "fwupdate_beta": RestUpdateDescription(
+        name="Beta Firmware Update",
+        key="fwupdate",
+        latest_version=lambda status: status["update"].get("beta_version"),
+        install=lambda wrapper: wrapper.async_trigger_ota_update(beta=True),
+        device_class=UpdateDeviceClass.FIRMWARE,
+        entity_category=EntityCategory.CONFIG,
+        entity_registry_enabled_default=False,
+    ),
+}
+
+RPC_UPDATES: Final = {
+    "fwupdate": RpcUpdateDescription(
+        name="Firmware Update",
+        key="sys",
+        sub_key="available_updates",
+        latest_version=lambda status: status.get("stable", {"version": None})[
+            "version"
+        ],
+        install=lambda wrapper: wrapper.async_trigger_ota_update(),
+        device_class=UpdateDeviceClass.FIRMWARE,
+        entity_category=EntityCategory.CONFIG,
+        entity_registry_enabled_default=True,
+    ),
+    "fwupdate_beta": RpcUpdateDescription(
+        name="Beta Firmware Update",
+        key="sys",
+        sub_key="available_updates",
+        latest_version=lambda status: status.get("beta", {"version": None})["version"],
+        install=lambda wrapper: wrapper.async_trigger_ota_update(beta=True),
+        device_class=UpdateDeviceClass.FIRMWARE,
+        entity_category=EntityCategory.CONFIG,
+        entity_registry_enabled_default=False,
+    ),
+}
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up update entities for Shelly component."""
+    if get_device_entry_gen(config_entry) == 2:
+        return async_setup_entry_rpc(
+            hass, config_entry, async_add_entities, RPC_UPDATES, RpcUpdateEntity
+        )
+
+    if not config_entry.data[CONF_SLEEP_PERIOD]:
+        async_setup_entry_rest(
+            hass,
+            config_entry,
+            async_add_entities,
+            REST_UPDATES,
+            RestUpdateEntity,
+        )
+
+
+class RestUpdateEntity(ShellyRestAttributeEntity, UpdateEntity):
+    """Represent a REST update entity."""
+
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
+    )
+    entity_description: RestUpdateDescription
+
+    def __init__(
+        self,
+        wrapper: BlockDeviceWrapper,
+        attribute: str,
+        description: RestEntityDescription,
+    ) -> None:
+        """Initialize update entity."""
+        super().__init__(wrapper, attribute, description)
+        self._in_progress_old_version: str | None = None
+
+    @property
+    def installed_version(self) -> str | None:
+        """Version currently in use."""
+        version = self.wrapper.device.status["update"]["old_version"]
+        if version is None:
+            return None
+
+        return cast(str, version)
+
+    @property
+    def latest_version(self) -> str | None:
+        """Latest version available for install."""
+        new_version = self.entity_description.latest_version(
+            self.wrapper.device.status,
+        )
+        if new_version is not None:
+            return cast(str, new_version)
+
+        return self.installed_version
+
+    @property
+    def in_progress(self) -> bool:
+        """Update installation in progress."""
+        return self._in_progress_old_version == self.installed_version
+
+    async def async_install(
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Install the latest firmware version."""
+        config_entry = self.wrapper.entry
+        block_wrapper = self.hass.data[DOMAIN][DATA_CONFIG_ENTRY][
+            config_entry.entry_id
+        ].get(BLOCK)
+        self._in_progress_old_version = self.installed_version
+        await self.entity_description.install(block_wrapper)
+
+
+class RpcUpdateEntity(ShellyRpcAttributeEntity, UpdateEntity):
+    """Represent a RPC update entity."""
+
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
+    )
+    entity_description: RpcUpdateDescription
+
+    def __init__(
+        self,
+        wrapper: RpcDeviceWrapper,
+        key: str,
+        attribute: str,
+        description: RpcEntityDescription,
+    ) -> None:
+        """Initialize update entity."""
+        super().__init__(wrapper, key, attribute, description)
+        self._in_progress_old_version: str | None = None
+
+    @property
+    def installed_version(self) -> str | None:
+        """Version currently in use."""
+        if self.wrapper.device.shelly is None:
+            return None
+
+        return cast(str, self.wrapper.device.shelly["ver"])
+
+    @property
+    def latest_version(self) -> str | None:
+        """Latest version available for install."""
+        new_version = self.entity_description.latest_version(
+            self.wrapper.device.status[self.key][self.entity_description.sub_key],
+        )
+        if new_version is not None:
+            return cast(str, new_version)
+
+        return self.installed_version
+
+    @property
+    def in_progress(self) -> bool:
+        """Update installation in progress."""
+        return self._in_progress_old_version == self.installed_version
+
+    async def async_install(
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Install the latest firmware version."""
+        self._in_progress_old_version = self.installed_version
+        await self.entity_description.install(self.wrapper)

tests/components/shelly/conftest.py

@@ -3,13 +3,21 @@ from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
-from homeassistant.components.shelly import BlockDeviceWrapper, RpcDeviceWrapper
+from homeassistant.components.shelly import (
+    BlockDeviceWrapper,
+    RpcDeviceWrapper,
+    RpcPollingWrapper,
+    ShellyDeviceRestWrapper,
+)
 from homeassistant.components.shelly.const import (
     BLOCK,
     DATA_CONFIG_ENTRY,
     DOMAIN,
     EVENT_SHELLY_CLICK,
+    REST,
+    REST_SENSORS_UPDATE_INTERVAL,
     RPC,
+    RPC_POLL,
 )
 from homeassistant.setup import async_setup_component
 
@@ -65,13 +73,27 @@ MOCK_CONFIG = {
     },
 }
 
-MOCK_SHELLY = {
+MOCK_SHELLY_COAP = {
     "mac": "test-mac",
     "auth": False,
     "fw": "20201124-092854/v1.9.0@57ac4ad8",
     "num_outputs": 2,
 }
 
+MOCK_SHELLY_RPC = {
+    "name": "Test Gen2",
+    "id": "shellyplus2pm-123456789abc",
+    "mac": "123456789ABC",
+    "model": "SNSW-002P16EU",
+    "gen": 2,
+    "fw_id": "20220830-130540/0.11.0-gfa1bc37",
+    "ver": "0.11.0",
+    "app": "Plus2PM",
+    "auth_en": False,
+    "auth_domain": None,
+    "profile": "cover",
+}
+
 MOCK_STATUS_COAP = {
     "update": {
         "status": "pending",
@@ -80,6 +102,7 @@ MOCK_STATUS_COAP = {
         "new_version": "some_new_version",
         "old_version": "some_old_version",
     },
+    "uptime": 5 * REST_SENSORS_UPDATE_INTERVAL,
 }
 
 
@@ -135,10 +158,11 @@ async def coap_wrapper(hass):
     device = Mock(
         blocks=MOCK_BLOCKS,
         settings=MOCK_SETTINGS,
-        shelly=MOCK_SHELLY,
+        shelly=MOCK_SHELLY_COAP,
         status=MOCK_STATUS_COAP,
         firmware_version="some fw string",
         update=AsyncMock(),
+        update_status=AsyncMock(),
         trigger_ota_update=AsyncMock(),
         trigger_reboot=AsyncMock(),
         initialized=True,
@@ -146,6 +170,10 @@ async def coap_wrapper(hass):
 
     hass.data[DOMAIN] = {DATA_CONFIG_ENTRY: {}}
     hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id] = {}
+    hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id][
+        REST
+    ] = ShellyDeviceRestWrapper(hass, device, config_entry)
+
     wrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id][
         BLOCK
     ] = BlockDeviceWrapper(hass, config_entry, device)
@@ -157,7 +185,7 @@ async def coap_wrapper(hass):
 
 @pytest.fixture
 async def rpc_wrapper(hass):
-    """Setups a coap wrapper with mocked device."""
+    """Setups a rpc wrapper with mocked device."""
     await async_setup_component(hass, "shelly", {})
 
     config_entry = MockConfigEntry(
@@ -171,7 +199,7 @@ async def rpc_wrapper(hass):
         call_rpc=AsyncMock(),
         config=MOCK_CONFIG,
         event={},
-        shelly=MOCK_SHELLY,
+        shelly=MOCK_SHELLY_RPC,
         status=MOCK_STATUS_RPC,
         firmware_version="some fw string",
         update=AsyncMock(),
@@ -183,10 +211,13 @@ async def rpc_wrapper(hass):
 
     hass.data[DOMAIN] = {DATA_CONFIG_ENTRY: {}}
     hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id] = {}
+    hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id][
+        RPC_POLL
+    ] = RpcPollingWrapper(hass, config_entry, device)
+
     wrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id][
         RPC
     ] = RpcDeviceWrapper(hass, config_entry, device)
-
     wrapper.async_setup()
 
     return wrapper

tests/components/shelly/test_button.py

@@ -7,15 +7,15 @@ from homeassistant.helpers.entity_registry import async_get
 
 
 async def test_block_button(hass: HomeAssistant, coap_wrapper):
-    """Test block device OTA button."""
+    """Test block device reboot button."""
     assert coap_wrapper
 
     entity_registry = async_get(hass)
     entity_registry.async_get_or_create(
         BUTTON_DOMAIN,
         DOMAIN,
-        "test_name_ota_update_beta",
+        "test_name_reboot",
-        suggested_object_id="test_name_ota_update_beta",
+        suggested_object_id="test_name_reboot",
         disabled_by=None,
     )
     hass.async_create_task(
@@ -23,37 +23,6 @@ async def test_block_button(hass: HomeAssistant, coap_wrapper):
     )
     await hass.async_block_till_done()
 
-    # stable channel button
-    state = hass.states.get("button.test_name_ota_update")
-    assert state
-    assert state.state == STATE_UNKNOWN
-
-    await hass.services.async_call(
-        BUTTON_DOMAIN,
-        SERVICE_PRESS,
-        {ATTR_ENTITY_ID: "button.test_name_ota_update"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    assert coap_wrapper.device.trigger_ota_update.call_count == 1
-    coap_wrapper.device.trigger_ota_update.assert_called_with(beta=False)
-
-    # beta channel button
-    state = hass.states.get("button.test_name_ota_update_beta")
-
-    assert state
-    assert state.state == STATE_UNKNOWN
-
-    await hass.services.async_call(
-        BUTTON_DOMAIN,
-        SERVICE_PRESS,
-        {ATTR_ENTITY_ID: "button.test_name_ota_update_beta"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    assert coap_wrapper.device.trigger_ota_update.call_count == 2
-    coap_wrapper.device.trigger_ota_update.assert_called_with(beta=True)
-
     # reboot button
     state = hass.states.get("button.test_name_reboot")
 
@@ -78,8 +47,8 @@ async def test_rpc_button(hass: HomeAssistant, rpc_wrapper):
     entity_registry.async_get_or_create(
         BUTTON_DOMAIN,
         DOMAIN,
-        "test_name_ota_update_beta",
+        "test_name_reboot",
-        suggested_object_id="test_name_ota_update_beta",
+        suggested_object_id="test_name_reboot",
         disabled_by=None,
     )
 
@@ -88,37 +57,6 @@ async def test_rpc_button(hass: HomeAssistant, rpc_wrapper):
     )
     await hass.async_block_till_done()
 
-    # stable channel button
-    state = hass.states.get("button.test_name_ota_update")
-    assert state
-    assert state.state == STATE_UNKNOWN
-
-    await hass.services.async_call(
-        BUTTON_DOMAIN,
-        SERVICE_PRESS,
-        {ATTR_ENTITY_ID: "button.test_name_ota_update"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    assert rpc_wrapper.device.trigger_ota_update.call_count == 1
-    rpc_wrapper.device.trigger_ota_update.assert_called_with(beta=False)
-
-    # beta channel button
-    state = hass.states.get("button.test_name_ota_update_beta")
-
-    assert state
-    assert state.state == STATE_UNKNOWN
-
-    await hass.services.async_call(
-        BUTTON_DOMAIN,
-        SERVICE_PRESS,
-        {ATTR_ENTITY_ID: "button.test_name_ota_update_beta"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    assert rpc_wrapper.device.trigger_ota_update.call_count == 2
-    rpc_wrapper.device.trigger_ota_update.assert_called_with(beta=True)
-
     # reboot button
     state = hass.states.get("button.test_name_reboot")
 

tests/components/shelly/test_diagnostics.py

@@ -6,6 +6,8 @@ from homeassistant.components.shelly.const import DOMAIN
 from homeassistant.components.shelly.diagnostics import TO_REDACT
 from homeassistant.core import HomeAssistant
 
+from .conftest import MOCK_STATUS_COAP
+
 from tests.components.diagnostics import get_diagnostics_for_config_entry
 
 RELAY_BLOCK_ID = 0
@@ -33,15 +35,7 @@ async def test_block_config_entry_diagnostics(
             "sw_version": coap_wrapper.sw_version,
         },
         "device_settings": {"coiot": {"update_period": 15}},
-        "device_status": {
+        "device_status": MOCK_STATUS_COAP,
-            "update": {
-                "beta_version": "some_beta_version",
-                "has_update": True,
-                "new_version": "some_new_version",
-                "old_version": "some_old_version",
-                "status": "pending",
-            }
-        },
     }
 
 

tests/components/shelly/test_update.py

@@ -0,0 +1,101 @@
+"""Tests for Shelly update platform."""
+from homeassistant.components.shelly.const import DOMAIN
+from homeassistant.components.update import DOMAIN as UPDATE_DOMAIN
+from homeassistant.components.update.const import SERVICE_INSTALL
+from homeassistant.const import ATTR_ENTITY_ID, STATE_ON, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_component import async_update_entity
+from homeassistant.helpers.entity_registry import async_get
+
+
+async def test_block_update(hass: HomeAssistant, coap_wrapper, monkeypatch):
+    """Test block device update entity."""
+    assert coap_wrapper
+
+    entity_registry = async_get(hass)
+    entity_registry.async_get_or_create(
+        UPDATE_DOMAIN,
+        DOMAIN,
+        "test_name_update",
+        suggested_object_id="test_name_update",
+        disabled_by=None,
+    )
+    hass.async_create_task(
+        hass.config_entries.async_forward_entry_setup(coap_wrapper.entry, UPDATE_DOMAIN)
+    )
+    await hass.async_block_till_done()
+
+    # update entity
+    await async_update_entity(hass, "update.test_name_firmware_update")
+    await hass.async_block_till_done()
+    state = hass.states.get("update.test_name_firmware_update")
+
+    assert state
+    assert state.state == STATE_ON
+
+    await hass.services.async_call(
+        UPDATE_DOMAIN,
+        SERVICE_INSTALL,
+        {ATTR_ENTITY_ID: "update.test_name_firmware_update"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert coap_wrapper.device.trigger_ota_update.call_count == 1
+
+    monkeypatch.setitem(coap_wrapper.device.status["update"], "old_version", None)
+    monkeypatch.setitem(coap_wrapper.device.status["update"], "new_version", None)
+
+    # update entity
+    await async_update_entity(hass, "update.test_name_firmware_update")
+    await hass.async_block_till_done()
+    state = hass.states.get("update.test_name_firmware_update")
+
+    assert state
+    assert state.state == STATE_UNKNOWN
+
+
+async def test_rpc_update(hass: HomeAssistant, rpc_wrapper, monkeypatch):
+    """Test rpc device update entity."""
+    assert rpc_wrapper
+
+    entity_registry = async_get(hass)
+    entity_registry.async_get_or_create(
+        UPDATE_DOMAIN,
+        DOMAIN,
+        "test_name_update",
+        suggested_object_id="test_name_update",
+        disabled_by=None,
+    )
+
+    hass.async_create_task(
+        hass.config_entries.async_forward_entry_setup(rpc_wrapper.entry, UPDATE_DOMAIN)
+    )
+    await hass.async_block_till_done()
+
+    # update entity
+    await async_update_entity(hass, "update.test_name_firmware_update")
+    await hass.async_block_till_done()
+    state = hass.states.get("update.test_name_firmware_update")
+
+    assert state
+    assert state.state == STATE_ON
+
+    await hass.services.async_call(
+        UPDATE_DOMAIN,
+        SERVICE_INSTALL,
+        {ATTR_ENTITY_ID: "update.test_name_firmware_update"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert rpc_wrapper.device.trigger_ota_update.call_count == 1
+
+    monkeypatch.setitem(rpc_wrapper.device.status["sys"], "available_updates", {})
+    rpc_wrapper.device.shelly = None
+
+    # update entity
+    await async_update_entity(hass, "update.test_name_firmware_update")
+    await hass.async_block_till_done()
+    state = hass.states.get("update.test_name_firmware_update")
+
+    assert state
+    assert state.state == STATE_UNKNOWN