Add button to trigger ota firmware update for Shelly devices

homeassistant/components/shelly/__init__.py

@@ -29,6 +29,7 @@
 
 from .const import (
     AIOSHELLY_DEVICE_TIMEOUT_SEC,
+    ATTR_BETA,
     ATTR_CHANNEL,
     ATTR_CLICK_TYPE,
     ATTR_DEVICE,
@@ -64,9 +65,16 @@
     get_rpc_device_name,
 )
 
-BLOCK_PLATFORMS: Final = ["binary_sensor", "cover", "light", "sensor", "switch"]
+BLOCK_PLATFORMS: Final = [
+    "binary_sensor",
+    "button",
+    "cover",
+    "light",
+    "sensor",
+    "switch",
+]
 BLOCK_SLEEPING_PLATFORMS: Final = ["binary_sensor", "sensor"]
-RPC_PLATFORMS: Final = ["binary_sensor", "light", "sensor", "switch"]
+RPC_PLATFORMS: Final = ["binary_sensor", "button", "light", "sensor", "switch"]
 _LOGGER: Final = logging.getLogger(__name__)
 
 COAP_SCHEMA: Final = vol.Schema(
@@ -411,6 +419,41 @@ def async_setup(self) -> None:
         self.device_id = entry.id
         self.device.subscribe_updates(self.async_set_updated_data)
 
+    async def async_trigger_ota_update(self, beta: bool = False) -> None:
+        """Trigger or schedule an ota update."""
+        update_data = self.device.status["update"]
+        _LOGGER.debug("OTA update service - update_data: %s", update_data)
+
+        if not update_data["has_update"] and not beta:
+            _LOGGER.warning("No OTA update available for device %s", self.name)
+            return
+
+        if beta and not update_data.get("beta_version"):
+            _LOGGER.warning(
+                "No OTA update on beta channel available for device %s", self.name
+            )
+            return
+
+        if update_data["status"] == "updating":
+            _LOGGER.warning("OTA update already in progress for %s", self.name)
+            return
+
+        new_version = update_data["new_version"]
+        if beta:
+            new_version = update_data["beta_version"]
+        _LOGGER.info(
+            "Start OTA update of device %s from '%s' to '%s'",
+            self.name,
+            self.device.firmware_version,
+            new_version,
+        )
+        try:
+            async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
+                result = await self.device.trigger_ota_update(beta=beta)
+        except (asyncio.TimeoutError, OSError) as err:
+            _LOGGER.exception("Error while perform ota update: %s", err)
+        _LOGGER.debug("Result of OTA update call: %s", result)
+
     def shutdown(self) -> None:
         """Shutdown the wrapper."""
         self.device.shutdown()
@@ -648,6 +691,42 @@ def async_setup(self) -> None:
         self.device_id = entry.id
         self.device.subscribe_updates(self.async_set_updated_data)
 
+    async def async_trigger_ota_update(self, beta: bool = False) -> None:
+        """Trigger an ota update."""
+
+        update_data = self.device.status["sys"]["available_updates"]
+        _LOGGER.debug("OTA update service - update_data: %s", update_data)
+
+        if not bool(update_data) or (not update_data.get("stable") and not beta):
+            _LOGGER.warning("No OTA update available for device %s", self.name)
+            return
+
+        if beta and not update_data.get(ATTR_BETA):
+            _LOGGER.warning(
+                "No OTA update on beta channel available for device %s", self.name
+            )
+            return
+
+        new_version = update_data.get("stable", {"version": ""})["version"]
+        if beta:
+            new_version = update_data.get(ATTR_BETA, {"version": ""})["version"]
+
+        assert self.device.shelly
+        _LOGGER.info(
+            "Start OTA update of device %s from '%s' to '%s'",
+            self.name,
+            self.device.firmware_version,
+            new_version,
+        )
+        result = None
+        try:
+            async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
+                result = await self.device.trigger_ota_update(beta=beta)
+        except (asyncio.TimeoutError, OSError) as err:
+            _LOGGER.exception("Error while perform ota update: %s", err)
+
+        _LOGGER.debug("Result of OTA update call: %s", result)
+
     async def shutdown(self) -> None:
         """Shutdown the wrapper."""
         await self.device.shutdown()

homeassistant/components/shelly/button.py

@@ -0,0 +1,71 @@
+"""Button for Shelly."""
+from __future__ import annotations
+
+from typing import cast
+
+from homeassistant.components.button import ButtonEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import ENTITY_CATEGORY_CONFIG
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.util import slugify
+
+from . import BlockDeviceWrapper, RpcDeviceWrapper
+from .const import BLOCK, CONF_OTA_BETA_CHANNEL, DATA_CONFIG_ENTRY, DOMAIN, RPC
+from .utils import get_block_device_name, get_device_entry_gen, get_rpc_device_name
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set buttons for device."""
+    wrapper: RpcDeviceWrapper | BlockDeviceWrapper | None = None
+    if get_device_entry_gen(config_entry) == 2:
+        if rpc_wrapper := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
+            config_entry.entry_id
+        ].get(RPC):
+            wrapper = cast(RpcDeviceWrapper, rpc_wrapper)
+    else:
+        if block_wrapper := hass.data[DOMAIN][DATA_CONFIG_ENTRY][
+            config_entry.entry_id
+        ].get(BLOCK):
+            wrapper = cast(BlockDeviceWrapper, block_wrapper)
+
+    if wrapper is not None:
+        async_add_entities([ShellyOtaUpdateButton(wrapper, config_entry)])
+
+
+class ShellyOtaUpdateButton(ButtonEntity):
+    """Defines a Shelly OTA update button."""
+
+    _attr_icon = "mdi:package-up"
+    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+
+    def __init__(
+        self, wrapper: RpcDeviceWrapper | BlockDeviceWrapper, entry: ConfigEntry
+    ) -> None:
+        """Initialize Shelly OTA update button."""
+        self._attr_device_info = DeviceInfo(
+            connections={(CONNECTION_NETWORK_MAC, wrapper.mac)}
+        )
+
+        if isinstance(wrapper, RpcDeviceWrapper):
+            device_name = get_rpc_device_name(wrapper.device)
+        else:
+            device_name = get_block_device_name(wrapper.device)
+
+        self._attr_name = f"{device_name} OTA Update"
+        self._attr_unique_id = slugify(self._attr_name)
+
+        self.entry = entry
+        self.wrapper = wrapper
+
+    async def async_press(self) -> None:
+        """Triggers the OTA update service."""
+        await self.wrapper.async_trigger_ota_update(
+            beta=bool(self.entry.options.get(CONF_OTA_BETA_CHANNEL))
+        )

homeassistant/components/shelly/const.py

@@ -90,6 +90,8 @@
 ATTR_DEVICE: Final = "device"
 ATTR_GENERATION: Final = "generation"
 CONF_SUBTYPE: Final = "subtype"
+ATTR_BETA: Final = "beta"
+CONF_OTA_BETA_CHANNEL: Final = "ota_beta_channel"
 
 BASIC_INPUTS_EVENTS_TYPES: Final = {"single", "long"}
 

homeassistant/components/shelly/switch.py

@@ -7,15 +7,21 @@
 
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import ENTITY_CATEGORY_CONFIG
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.util import slugify
 
 from . import BlockDeviceWrapper, RpcDeviceWrapper
-from .const import BLOCK, DATA_CONFIG_ENTRY, DOMAIN, RPC
+from .const import BLOCK, CONF_OTA_BETA_CHANNEL, DATA_CONFIG_ENTRY, DOMAIN, RPC
 from .entity import ShellyBlockEntity, ShellyRpcEntity
 from .utils import (
     async_remove_shelly_entity,
+    get_block_device_name,
     get_device_entry_gen,
+    get_rpc_device_name,
     get_rpc_key_ids,
     is_block_channel_type_light,
     is_rpc_channel_type_light,
@@ -42,6 +48,8 @@ async def async_setup_block_entry(
     """Set up entities for block device."""
     wrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id][BLOCK]
 
+    async_add_entities([ShellyOtaUpdateBetaChannelSwitch(wrapper, config_entry)])
+
     # In roller mode the relay blocks exist but do not contain required info
     if (
         wrapper.model in ["SHSW-21", "SHSW-25"]
@@ -74,6 +82,9 @@ async def async_setup_rpc_entry(
 ) -> None:
     """Set up entities for RPC device."""
     wrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id][RPC]
+
+    async_add_entities([ShellyOtaUpdateBetaChannelSwitch(wrapper, config_entry)])
+
     switch_key_ids = get_rpc_key_ids(wrapper.device.status, "switch")
 
     switch_ids = []
@@ -144,3 +155,46 @@ async def async_turn_on(self, **kwargs: Any) -> None:
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off relay."""
         await self.call_rpc("Switch.Set", {"id": self._id, "on": False})
+
+
+class ShellyOtaUpdateBetaChannelSwitch(SwitchEntity):
+    """Defines a Shelly OTA update beta channel switch."""
+
+    _attr_icon = "mdi:flask-outline"
+    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+
+    def __init__(
+        self, wrapper: RpcDeviceWrapper | BlockDeviceWrapper, entry: ConfigEntry
+    ) -> None:
+        """Initialize Shelly OTA update beta channel switch."""
+        self._attr_device_info = DeviceInfo(
+            connections={(CONNECTION_NETWORK_MAC, wrapper.mac)}
+        )
+
+        if isinstance(wrapper, RpcDeviceWrapper):
+            device_name = get_rpc_device_name(wrapper.device)
+        else:
+            device_name = get_block_device_name(wrapper.device)
+
+        self._attr_name = f"{device_name} OTA Update Beta channel"
+        self._attr_unique_id = slugify(self._attr_name)
+
+        self.entry = entry
+        self.wrapper = wrapper
+
+    @property
+    def is_on(self) -> bool:
+        """Return true if device is on."""
+        return bool(self.entry.options.get(CONF_OTA_BETA_CHANNEL))
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the device on."""
+        self.hass.config_entries.async_update_entry(
+            self.entry, options={**self.entry.options, CONF_OTA_BETA_CHANNEL: True}
+        )
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the device off."""
+        self.hass.config_entries.async_update_entry(
+            self.entry, options={**self.entry.options, CONF_OTA_BETA_CHANNEL: False}
+        )

tests/components/shelly/conftest.py

@@ -71,8 +71,25 @@
     "num_outputs": 2,
 }
 
-MOCK_STATUS = {
+MOCK_STATUS_COAP = {
+    "update": {
+        "status": "pending",
+        "has_update": True,
+        "beta_version": "some_beta_version",
+        "new_version": "some_new_version",
+        "old_version": "some_old_version",
+    },
+}
+
+
+MOCK_STATUS_RPC = {
     "switch:0": {"output": True},
+    "sys": {
+        "available_updates": {
+            "beta": {"version": "some_beta_version"},
+            "stable": {"version": "some_beta_version"},
+        }
+    },
 }
 
 
@@ -117,8 +134,10 @@ async def coap_wrapper(hass):
         blocks=MOCK_BLOCKS,
         settings=MOCK_SETTINGS,
         shelly=MOCK_SHELLY,
+        status=MOCK_STATUS_COAP,
         firmware_version="some fw string",
         update=AsyncMock(),
+        trigger_ota_update=AsyncMock(),
         initialized=True,
     )
 
@@ -150,9 +169,10 @@ async def rpc_wrapper(hass):
         config=MOCK_CONFIG,
         event={},
         shelly=MOCK_SHELLY,
-        status=MOCK_STATUS,
+        status=MOCK_STATUS_RPC,
         firmware_version="some fw string",
         update=AsyncMock(),
+        trigger_ota_update=AsyncMock(),
         initialized=True,
         shutdown=AsyncMock(),
     )