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,
@@ -56,6 +57,7 @@
     SLEEP_PERIOD_MULTIPLIER,
     UPDATE_PERIOD_MULTIPLIER,
 )
+from .service import async_services_setup
 from .utils import (
     get_block_device_name,
     get_block_device_sleep_period,
@@ -64,9 +66,16 @@
     get_rpc_device_name,
 )
 
-BLOCK_PLATFORMS: Final = ["binary_sensor", "cover", "light", "sensor", "switch"]
-BLOCK_SLEEPING_PLATFORMS: Final = ["binary_sensor", "sensor"]
-RPC_PLATFORMS: Final = ["binary_sensor", "light", "sensor", "switch"]
+BLOCK_PLATFORMS: Final = [
+    "binary_sensor",
+    "button",
+    "cover",
+    "light",
+    "sensor",
+    "switch",
+]
+BLOCK_SLEEPING_PLATFORMS: Final = ["binary_sensor", "button", "sensor"]
+RPC_PLATFORMS: Final = ["binary_sensor", "button", "light", "sensor", "switch"]
 _LOGGER: Final = logging.getLogger(__name__)
 
 COAP_SCHEMA: Final = vol.Schema(
@@ -182,6 +191,8 @@ def _async_device_online(_: Any) -> None:
         _LOGGER.debug("Setting up offline block device %s", entry.title)
         await async_block_device_setup(hass, entry, device)
 
+    await async_services_setup(hass)
+
     return True
 
 
@@ -229,6 +240,8 @@ async def async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool
 
     hass.config_entries.async_setup_platforms(entry, RPC_PLATFORMS)
 
+    await async_services_setup(hass)
+
     return True
 
 
@@ -273,6 +286,9 @@ def __init__(
         self._last_mode: str | None = None
         self._last_effect: int | None = None
 
+        self._ota_update_pending = False
+        self._ota_update_params = {ATTR_BETA: False}
+
         entry.async_on_unload(
             self.async_add_listener(self._async_device_updates_handler)
         )
@@ -306,6 +322,13 @@ def _async_device_updates_handler(self) -> None:
 
                 break
 
+        # check if OTA update is scheduled
+        if self._ota_update_pending:
+            self._ota_update_pending = False
+            self.hass.async_create_task(
+                self.async_trigger_ota_update(**self._ota_update_params)
+            )
+
         # Check for input events and config change
         cfg_changed = 0
         for block in self.device.blocks:
@@ -411,6 +434,63 @@ 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."""
+
+        async def _async_do_ota_update() -> None:
+            """Perform an ota update."""
+            beta_channel = self._ota_update_params[ATTR_BETA]
+            update_data = self.device.status["update"]
+            new_version = update_data["new_version"]
+            if beta_channel:
+                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,
+            )
+            result = None
+            try:
+                async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
+                    result = await self.device.trigger_ota_update(beta=beta_channel)
+            except (asyncio.TimeoutError, OSError) as err:
+                _LOGGER.exception("Error while perform ota update: %s", err)
+
+            _LOGGER.debug("Result of OTA update call: %s", result)
+
+        update_data = self.device.status["update"]
+        _LOGGER.debug("OTA update service - update_data: %s", update_data)
+
+        if self._ota_update_pending:
+            _LOGGER.warning(
+                "OTA update already scheduled for sleeping device %s", self.name
+            )
+            return
+
+        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
+
+        if self.entry.data.get(CONF_SLEEP_PERIOD):
+            self._ota_update_pending = True
+            self._ota_update_params = {ATTR_BETA: beta}
+            _LOGGER.info("OTA update scheduled for sleeping device %s", self.name)
+        else:
+            self._ota_update_params = {ATTR_BETA: beta}
+            await _async_do_ota_update()
+
     def shutdown(self) -> None:
         """Shutdown the wrapper."""
         self.device.shutdown()
@@ -549,6 +629,8 @@ def __init__(
         self.entry = entry
         self.device = device
 
+        self._ota_update_params = {ATTR_BETA: False}
+
         self._debounced_reload = Debouncer(
             hass,
             _LOGGER,
@@ -648,6 +730,49 @@ 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."""
+
+        async def _async_do_ota_update() -> None:
+            """Perform an ota update."""
+            beta_channel = self._ota_update_params[ATTR_BETA]
+            update_data = self.device.status["sys"]["available_updates"]
+            new_version = update_data.get("stable", {"version": ""})["version"]
+            if beta_channel:
+                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_channel)
+            except (asyncio.TimeoutError, OSError) as err:
+                _LOGGER.exception("Error while perform ota update: %s", err)
+
+            _LOGGER.debug("Result of OTA update call: %s", result)
+
+        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
+
+        self._ota_update_params = {ATTR_BETA: beta}
+        await _async_do_ota_update()
+
     async def shutdown(self) -> None:
         """Shutdown the wrapper."""
         await self.device.shutdown()

homeassistant/components/shelly/button.py

@@ -0,0 +1,84 @@
+"""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 ATTR_DEVICE_ID, 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 (
+    ATTR_BETA,
+    BLOCK,
+    CONF_OTA_BETA_CHANNEL,
+    DATA_CONFIG_ENTRY,
+    DOMAIN,
+    RPC,
+    SERVICE_OTA_UPDATE,
+)
+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.hass.services.async_call(
+            DOMAIN,
+            SERVICE_OTA_UPDATE,
+            {
+                ATTR_DEVICE_ID: self.wrapper.device_id,
+                ATTR_BETA: self.entry.options.get(CONF_OTA_BETA_CHANNEL),
+            },
+        )

homeassistant/components/shelly/const.py

@@ -10,6 +10,8 @@
 DOMAIN: Final = "shelly"
 REST: Final = "rest"
 RPC: Final = "rpc"
+SERVICE_OTA_UPDATE: Final = "ota_update"
+SERVICES: Final = [SERVICE_OTA_UPDATE]
 
 CONF_COAP_PORT: Final = "coap_port"
 DEFAULT_COAP_PORT: Final = 5683
@@ -90,6 +92,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/service.py

@@ -0,0 +1,36 @@
+"""Services for the Shelly integration."""
+from __future__ import annotations
+
+from homeassistant.const import ATTR_AREA_ID, ATTR_DEVICE_ID
+from homeassistant.core import HomeAssistant, ServiceCall
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.service import async_extract_config_entry_ids
+
+from .const import ATTR_BETA, BLOCK, DATA_CONFIG_ENTRY, DOMAIN, RPC, SERVICE_OTA_UPDATE
+from .utils import get_device_entry_gen
+
+
+async def async_services_setup(hass: HomeAssistant) -> None:
+    """Set up services."""
+
+    async def async_service_ota_update(call: ServiceCall) -> None:
+        """Trigger OTA update."""
+        if not (call.data.get(ATTR_DEVICE_ID) or call.data.get(ATTR_AREA_ID)):
+            raise HomeAssistantError("No target selected for OTA update")
+
+        beta_channel = bool(call.data.get(ATTR_BETA))
+
+        entry_ids = await async_extract_config_entry_ids(hass, call)
+        for entry_id in entry_ids:
+            entry = hass.config_entries.async_get_entry(entry_id)
+            if not (entry and entry.domain == DOMAIN):
+                continue
+
+            if active_entry := hass.data[DOMAIN][DATA_CONFIG_ENTRY].get(entry.entry_id):
+                if get_device_entry_gen(entry) == 2:
+                    wrapper = active_entry.get(RPC)
+                else:
+                    wrapper = active_entry.get(BLOCK)
+                await wrapper.async_trigger_ota_update(beta=beta_channel)
+
+    hass.services.async_register(DOMAIN, SERVICE_OTA_UPDATE, async_service_ota_update)

homeassistant/components/shelly/services.yaml

@@ -0,0 +1,19 @@
+# shelly service descriptions.
+
+ota_update:
+  name: OTA Update
+  description: Trigger an over-the-air (OTA) update.
+  target:
+    device:
+      integration: shelly
+    entity:
+      integration: none
+  fields:
+    beta:
+      name: Beta
+      description: Run firmware update to beta version (if available)
+      required: false
+      default: false
+      example: true
+      selector:
+        boolean: