Beta firmware update switch for Connect integrations (#155370)

Author: puddly

homeassistant/components/homeassistant_connect_zbt2/__init__.py

@@ -2,20 +2,35 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 import logging
 import os.path
 
+from homeassistant.components.homeassistant_hardware.coordinator import (
+    FirmwareUpdateCoordinator,
+)
 from homeassistant.components.usb import USBDevice, async_register_port_event_callback
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.typing import ConfigType
 
-from .const import DEVICE, DOMAIN
+from .const import DEVICE, DOMAIN, NABU_CASA_FIRMWARE_RELEASES_URL
 
 _LOGGER = logging.getLogger(__name__)
 
+type HomeAssistantConnectZBT2ConfigEntry = ConfigEntry[HomeAssistantConnectZBT2Data]
+
+
+@dataclass
+class HomeAssistantConnectZBT2Data:
+    """Runtime data definition."""
+
+    coordinator: FirmwareUpdateCoordinator
+
+
 CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)
 
 
@@ -49,7 +64,9 @@ def async_port_event_callback(
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, entry: HomeAssistantConnectZBT2ConfigEntry
+) -> bool:
     """Set up a Home Assistant Connect ZBT-2 config entry."""
 
     # Postpone loading the config entry if the device is missing
@@ -60,12 +77,23 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             translation_key="device_disconnected",
         )
 
-    await hass.config_entries.async_forward_entry_setups(entry, ["update"])
+    # Create and store the firmware update coordinator in runtime_data
+    session = async_get_clientsession(hass)
+    coordinator = FirmwareUpdateCoordinator(
+        hass,
+        entry,
+        session,
+        NABU_CASA_FIRMWARE_RELEASES_URL,
+    )
+    entry.runtime_data = HomeAssistantConnectZBT2Data(coordinator)
+
+    await hass.config_entries.async_forward_entry_setups(entry, ["switch", "update"])
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, entry: HomeAssistantConnectZBT2ConfigEntry
+) -> bool:
     """Unload a config entry."""
-    await hass.config_entries.async_unload_platforms(entry, ["update"])
-    return True
+    return await hass.config_entries.async_unload_platforms(entry, ["switch", "update"])

homeassistant/components/homeassistant_connect_zbt2/const.py

@@ -3,7 +3,7 @@
 DOMAIN = "homeassistant_connect_zbt2"
 
 NABU_CASA_FIRMWARE_RELEASES_URL = (
-    "https://api.github.com/repos/NabuCasa/silabs-firmware-builder/releases/latest"
+    "https://api.github.com/repos/NabuCasa/silabs-firmware-builder/releases"
 )
 
 FIRMWARE = "firmware"

homeassistant/components/homeassistant_connect_zbt2/icons.json

@@ -0,0 +1,9 @@
+{
+  "entity": {
+    "switch": {
+      "beta_firmware": {
+        "default": "mdi:test-tube"
+      }
+    }
+  }
+}

homeassistant/components/homeassistant_connect_zbt2/strings.json

@@ -90,6 +90,13 @@
       }
     }
   },
+  "entity": {
+    "switch": {
+      "beta_firmware": {
+        "name": "Beta firmware updates"
+      }
+    }
+  },
   "exceptions": {
     "device_disconnected": {
       "message": "The device is not plugged in"

homeassistant/components/homeassistant_connect_zbt2/switch.py

@@ -0,0 +1,54 @@
+"""Home Assistant Connect ZBT-2 switch entities."""
+
+from __future__ import annotations
+
+import logging
+
+from homeassistant.components.homeassistant_hardware.coordinator import (
+    FirmwareUpdateCoordinator,
+)
+from homeassistant.components.homeassistant_hardware.switch import (
+    BaseBetaFirmwareSwitch,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+
+from . import HomeAssistantConnectZBT2ConfigEntry
+from .const import DOMAIN, HARDWARE_NAME, SERIAL_NUMBER
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: HomeAssistantConnectZBT2ConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
+    """Set up the switch platform for Home Assistant Connect ZBT-2."""
+    async_add_entities(
+        [BetaFirmwareSwitch(config_entry.runtime_data.coordinator, config_entry)]
+    )
+
+
+class BetaFirmwareSwitch(BaseBetaFirmwareSwitch):
+    """Home Assistant Connect ZBT-2 beta firmware switch."""
+
+    def __init__(
+        self,
+        coordinator: FirmwareUpdateCoordinator,
+        config_entry: HomeAssistantConnectZBT2ConfigEntry,
+    ) -> None:
+        """Initialize the beta firmware switch."""
+        super().__init__(coordinator, config_entry)
+
+        serial_number = self._config_entry.data[SERIAL_NUMBER]
+
+        self._attr_unique_id = f"{serial_number}_beta_firmware"
+        self._attr_device_info = DeviceInfo(
+            identifiers={(DOMAIN, serial_number)},
+            name=f"{HARDWARE_NAME} ({serial_number})",
+            model=HARDWARE_NAME,
+            manufacturer="Nabu Casa",
+            serial_number=serial_number,
+        )

homeassistant/components/homeassistant_connect_zbt2/update.py

@@ -4,8 +4,6 @@
 
 import logging
 
-import aiohttp
-
 from homeassistant.components.homeassistant_hardware.coordinator import (
     FirmwareUpdateCoordinator,
 )
@@ -19,22 +17,14 @@
     ResetTarget,
 )
 from homeassistant.components.update import UpdateDeviceClass
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr, entity_registry as er
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import (
-    DOMAIN,
-    FIRMWARE,
-    FIRMWARE_VERSION,
-    HARDWARE_NAME,
-    NABU_CASA_FIRMWARE_RELEASES_URL,
-    SERIAL_NUMBER,
-)
+from . import HomeAssistantConnectZBT2ConfigEntry
+from .const import DOMAIN, FIRMWARE, FIRMWARE_VERSION, HARDWARE_NAME, SERIAL_NUMBER
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -91,8 +81,7 @@
 
 def _async_create_update_entity(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
-    session: aiohttp.ClientSession,
+    config_entry: HomeAssistantConnectZBT2ConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> FirmwareUpdateEntity:
     """Create an update entity that handles firmware type changes."""
@@ -111,12 +100,7 @@ def _async_create_update_entity(
     entity = FirmwareUpdateEntity(
         device=config_entry.data["device"],
         config_entry=config_entry,
-        update_coordinator=FirmwareUpdateCoordinator(
-            hass,
-            config_entry,
-            session,
-            NABU_CASA_FIRMWARE_RELEASES_URL,
-        ),
+        update_coordinator=config_entry.runtime_data.coordinator,
         entity_description=entity_description,
     )
 
@@ -126,11 +110,7 @@ def firmware_type_changed(
         """Replace the current entity when the firmware type changes."""
         er.async_get(hass).async_remove(entity.entity_id)
         async_add_entities(
-            [
-                _async_create_update_entity(
-                    hass, config_entry, session, async_add_entities
-                )
-            ]
+            [_async_create_update_entity(hass, config_entry, async_add_entities)]
         )
 
     entity.async_on_remove(
@@ -142,14 +122,11 @@ def firmware_type_changed(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HomeAssistantConnectZBT2ConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the firmware update config entry."""
-    session = async_get_clientsession(hass)
-    entity = _async_create_update_entity(
-        hass, config_entry, session, async_add_entities
-    )
+    entity = _async_create_update_entity(hass, config_entry, async_add_entities)
 
     async_add_entities([entity])
 
@@ -162,7 +139,7 @@ class FirmwareUpdateEntity(BaseFirmwareUpdateEntity):
     def __init__(
         self,
         device: str,
-        config_entry: ConfigEntry,
+        config_entry: HomeAssistantConnectZBT2ConfigEntry,
         update_coordinator: FirmwareUpdateCoordinator,
         entity_description: FirmwareUpdateEntityDescription,
     ) -> None:

homeassistant/components/homeassistant_hardware/manifest.json

@@ -8,6 +8,6 @@
   "integration_type": "system",
   "requirements": [
     "universal-silabs-flasher==0.0.35",
-    "ha-silabs-firmware-client==0.2.0"
+    "ha-silabs-firmware-client==0.3.0"
   ]
 }

homeassistant/components/homeassistant_hardware/switch.py

@@ -0,0 +1,64 @@
+"""Home Assistant Hardware base beta firmware switch entity."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from homeassistant.components.switch import SwitchEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import EntityCategory
+from homeassistant.helpers.restore_state import RestoreEntity
+
+from .coordinator import FirmwareUpdateCoordinator
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class BaseBetaFirmwareSwitch(SwitchEntity, RestoreEntity):
+    """Base switch to enable beta firmware updates."""
+
+    _attr_has_entity_name = True
+    _attr_entity_category = EntityCategory.CONFIG
+    _attr_entity_registry_enabled_default = False
+    _attr_translation_key = "beta_firmware"
+
+    def __init__(
+        self,
+        coordinator: FirmwareUpdateCoordinator,
+        config_entry: ConfigEntry,
+    ) -> None:
+        """Initialize the beta firmware switch."""
+        self._coordinator = coordinator
+        self._config_entry = config_entry
+
+    async def async_added_to_hass(self) -> None:
+        """Handle entity which will be added to hass."""
+        await super().async_added_to_hass()
+
+        # Restore the last state
+        last_state = await self.async_get_last_state()
+        if last_state is not None:
+            self._attr_is_on = last_state.state == "on"
+        else:
+            self._attr_is_on = False
+
+        # Apply the restored state to the coordinator
+        await self._update_coordinator_prerelease()
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn on beta firmware updates."""
+        self._attr_is_on = True
+        self.async_write_ha_state()
+        await self._update_coordinator_prerelease()
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn off beta firmware updates."""
+        self._attr_is_on = False
+        self.async_write_ha_state()
+        await self._update_coordinator_prerelease()
+
+    async def _update_coordinator_prerelease(self) -> None:
+        """Update the coordinator with the current prerelease setting."""
+        self._coordinator.client.update_prerelease(bool(self._attr_is_on))
+        await self._coordinator.async_refresh()