Keep the device firmware version in sync (#427)

* Sync the device firmware version to its update entity

* Add a unit test

* Add event deduplication

* Shuffle code around a little

* Rename `DeviceUpdatedEvent` to `DeviceFirmwareInfoUpdatedEvent`

Author: puddly

tests/test_device.py

@@ -1,6 +1,7 @@
 """Test ZHA device switch."""
 
 import asyncio
+from datetime import UTC, datetime
 import logging
 import time
 from unittest import mock
@@ -13,7 +14,7 @@
 from zigpy.quirks.v2 import DeviceAlertLevel, DeviceAlertMetadata, QuirkBuilder
 import zigpy.types
 from zigpy.zcl.clusters import general
-from zigpy.zcl.clusters.general import PowerConfiguration
+from zigpy.zcl.clusters.general import Ota, PowerConfiguration
 from zigpy.zcl.foundation import Status, WriteAttributesResponse
 import zigpy.zdo.types as zdo_t
 
@@ -43,7 +44,11 @@
 from zha.application.platforms.sensor import LQISensor, RSSISensor
 from zha.application.platforms.switch import Switch
 from zha.exceptions import ZHAException
-from zha.zigbee.device import ClusterBinding, get_device_automation_triggers
+from zha.zigbee.device import (
+    ClusterBinding,
+    DeviceFirmwareInfoUpdatedEvent,
+    get_device_automation_triggers,
+)
 from zha.zigbee.group import Group
 
 
@@ -742,7 +747,7 @@ async def test_device_properties(
 
     assert zha_device.power_configuration_ch is None
     assert zha_device.basic_ch is not None
-    assert zha_device.sw_version is None
+    assert zha_device.firmware_version is None
 
     assert len(zha_device.platform_entities) == 3
 
@@ -781,6 +786,50 @@ async def test_device_properties(
     assert zha_device.is_coordinator is None
 
 
+async def test_device_firmware_version_syncing(zha_gateway: Gateway) -> None:
+    """Test device firmware version syncing."""
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/philips-sml001.json",
+    )
+
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    # Register a callback to listen for device updates
+    update_callback = mock.Mock()
+    zha_device.on_event(DeviceFirmwareInfoUpdatedEvent.event_type, update_callback)
+
+    # The firmware version is restored on device initialization
+    assert zha_device.firmware_version == "0x42006bb7"
+
+    # If we update the entity, the device updates as well
+    update_entity = get_entity(zha_device, platform=Platform.UPDATE)
+    update_entity._ota_cluster_handler.attribute_updated(
+        attrid=Ota.AttributeDefs.current_file_version.id,
+        value=zigpy.types.uint32_t(0xABCD1234),
+        timestamp=datetime.now(UTC),
+    )
+
+    assert zha_device.firmware_version == "0xabcd1234"
+
+    # Duplicate updates are ignored
+    update_entity._ota_cluster_handler.attribute_updated(
+        attrid=Ota.AttributeDefs.current_file_version.id,
+        value=zigpy.types.uint32_t(0xABCD1234),
+        timestamp=datetime.now(UTC),
+    )
+
+    assert zha_device.firmware_version == "0xabcd1234"
+    assert update_callback.mock_calls == [
+        call(
+            DeviceFirmwareInfoUpdatedEvent(
+                old_firmware_version="0x42006bb7",
+                new_firmware_version="0xabcd1234",
+            )
+        )
+    ]
+
+
 async def test_quirks_v2_device_renaming(zha_gateway: Gateway) -> None:
     """Test quirks v2 device renaming."""
     registry = DeviceRegistry()

tests/test_update.py

@@ -153,7 +153,6 @@ async def setup_test_data(
     )
 
     zha_device = await join_zigpy_device(zha_gateway, zigpy_device)
-    zha_device.async_update_sw_build_id(installed_fw_version)
 
     return zha_device, ota_cluster, fw_image, installed_fw_version
 

zha/application/const.py

@@ -192,6 +192,7 @@ def pretty_name(self) -> str:
 ZHA_CLUSTER_HANDLER_CFG_DONE = "zha_channel_cfg_done"
 ZHA_CLUSTER_HANDLER_READS_PER_REQ = 5
 ZHA_EVENT = "zha_event"
+ZHA_DEVICE_UPDATED_EVENT = "zha_device_updated_event"
 ZHA_GW_MSG = "zha_gateway_message"
 ZHA_GW_MSG_DEVICE_FULL_INIT = "device_fully_initialized"
 ZHA_GW_MSG_DEVICE_INFO = "device_info"

zha/application/platforms/update.py

@@ -7,7 +7,7 @@
 import functools
 import itertools
 import logging
-from typing import TYPE_CHECKING, Any, Final, final
+from typing import TYPE_CHECKING, Any, Final
 
 from zigpy.ota import OtaImagesResult, OtaImageWithMetadata
 from zigpy.zcl.clusters.general import Ota, QueryNextImageCommand
@@ -52,7 +52,6 @@ class UpdateEntityFeature(IntFlag):
     RELEASE_NOTES = 16
 
 
-ATTR_BACKUP: Final = "backup"
 ATTR_INSTALLED_VERSION: Final = "installed_version"
 ATTR_IN_PROGRESS: Final = "in_progress"
 ATTR_UPDATE_PERCENTAGE: Final = "update_percentage"
@@ -69,15 +68,13 @@ class UpdateEntityInfo(BaseEntityInfo):
 
     supported_features: UpdateEntityFeature
     device_class: UpdateDeviceClass
-    entity_category: EntityCategory
 
 
 class BaseFirmwareUpdateEntity(PlatformEntity):
     """Base representation of a ZHA firmware update entity."""
 
     PLATFORM = Platform.UPDATE
 
-    _unique_id_suffix = "firmware_update"
     _attr_entity_category = EntityCategory.CONFIG
     _attr_device_class = UpdateDeviceClass.FIRMWARE
     _attr_supported_features = (
@@ -105,7 +102,16 @@ def info_object(self) -> UpdateEntityInfo:
     def state(self):
         """Get the state for the entity."""
         response = super().state
-        response.update(self.state_attributes)
+        if (release_summary := self.release_summary) is not None:
+            release_summary = release_summary[:255]
+
+        response[ATTR_INSTALLED_VERSION] = self.installed_version
+        response[ATTR_IN_PROGRESS] = self.in_progress
+        response[ATTR_UPDATE_PERCENTAGE] = self.update_percentage
+        response[ATTR_LATEST_VERSION] = self.latest_version
+        response[ATTR_RELEASE_SUMMARY] = release_summary
+        response[ATTR_RELEASE_NOTES] = self.release_notes
+        response[ATTR_RELEASE_URL] = self.release_url
         return response
 
     @property
@@ -161,23 +167,6 @@ def supported_features(self) -> UpdateEntityFeature:
         """Flag supported features."""
         return self._attr_supported_features
 
-    @final
-    @property
-    def state_attributes(self) -> dict[str, Any] | None:
-        """Return state attributes."""
-        if (release_summary := self.release_summary) is not None:
-            release_summary = release_summary[:255]
-
-        return {
-            ATTR_INSTALLED_VERSION: self.installed_version,
-            ATTR_IN_PROGRESS: self.in_progress,
-            ATTR_UPDATE_PERCENTAGE: self.update_percentage,
-            ATTR_LATEST_VERSION: self.latest_version,
-            ATTR_RELEASE_SUMMARY: release_summary,
-            ATTR_RELEASE_NOTES: self.release_notes,
-            ATTR_RELEASE_URL: self.release_url,
-        }
-
     def _get_cluster_version(self) -> str | None:
         """Synchronize current file version with the cluster."""
 
@@ -303,6 +292,8 @@ async def on_remove(self) -> None:
 class FirmwareUpdateEntity(BaseFirmwareUpdateEntity):
     """Representation of a ZHA firmware update entity."""
 
+    _unique_id_suffix = "firmware_update"
+
     def __init__(
         self,
         cluster_handlers: list[ClusterHandler],
@@ -326,6 +317,8 @@ def __init__(
 class FirmwareUpdateServerEntity(BaseFirmwareUpdateEntity):
     """Representation of a ZHA firmware update entity."""
 
+    _unique_id_suffix = "firmware_update"
+
     def __init__(
         self,
         cluster_handlers: list[ClusterHandler],

zha/zigbee/cluster_handlers/general.py

@@ -691,7 +691,6 @@ def cluster_command(
             self.cluster.update_attribute(
                 Ota.AttributeDefs.current_file_version.id, current_file_version
             )
-            self._endpoint.device.sw_version = current_file_version
 
 
 @registries.CLUSTER_HANDLER_REGISTRY.register(Partition.cluster_id)

zha/zigbee/device.py

@@ -6,7 +6,7 @@
 
 import asyncio
 from collections import defaultdict
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 import copy
 import dataclasses
 from dataclasses import dataclass
@@ -61,10 +61,16 @@
     UNKNOWN_MODEL,
     ZHA_CLUSTER_HANDLER_CFG_DONE,
     ZHA_CLUSTER_HANDLER_MSG,
+    ZHA_DEVICE_UPDATED_EVENT,
     ZHA_EVENT,
 )
 from zha.application.helpers import convert_to_zcl_values, convert_zcl_value
-from zha.application.platforms import BaseEntityInfo, PlatformEntity
+from zha.application.platforms import (
+    BaseEntityInfo,
+    EntityStateChangedEvent,
+    PlatformEntity,
+)
+from zha.const import STATE_CHANGED
 from zha.event import EventBase
 from zha.exceptions import ZHAException
 from zha.mixins import LogMixin
@@ -147,6 +153,17 @@ class ZHAEvent:
     event: Final[str] = ZHA_EVENT
 
 
+@dataclass(kw_only=True, frozen=True)
+class DeviceFirmwareInfoUpdatedEvent:
+    """Event generated when the device firmware information has changed."""
+
+    event_type: Final[str] = ZHA_DEVICE_UPDATED_EVENT
+    event: Final[str] = ZHA_DEVICE_UPDATED_EVENT
+
+    old_firmware_version: str | None
+    new_firmware_version: str | None
+
+
 @dataclass(kw_only=True, frozen=True)
 class ClusterHandlerConfigurationComplete:
     """Event generated when all cluster handlers are configured."""
@@ -252,7 +269,7 @@ def __init__(
         self._power_config_ch: ClusterHandler | None = None
         self._identify_ch: ClusterHandler | None = None
         self._basic_ch: ClusterHandler | None = None
-        self._sw_build_id: int | None = None
+        self._firmware_version: str | None = None
 
         device_options = _gateway.config.config.device_options
         if self.is_mains_powered:
@@ -272,15 +289,20 @@ def __init__(
         self._pending_entities: list[PlatformEntity] = []
         self.semaphore: asyncio.Semaphore = asyncio.Semaphore(3)
 
+        self._on_remove_callbacks: list[Callable[[], None]] = []
+
         self._zdo_handler: ZDOClusterHandler = ZDOClusterHandler(self)
         self._zdo_handler.on_add()
+        self._on_remove_callbacks.append(self._zdo_handler.on_remove)
 
         self.status: DeviceStatus = DeviceStatus.CREATED
 
         self._endpoints: dict[int, Endpoint] = {}
         for ep_id, endpoint in zigpy_device.endpoints.items():
             if ep_id != 0:
-                self._endpoints[ep_id] = Endpoint.new(endpoint, self)
+                ep = Endpoint.new(endpoint, self)
+                self._endpoints[ep_id] = ep
+                self._on_remove_callbacks.append(ep.on_remove)
 
     def __repr__(self) -> str:
         """Return a string representation of the device."""
@@ -557,14 +579,9 @@ def zigbee_signature(self) -> dict[str, Any]:
         }
 
     @property
-    def sw_version(self) -> int | None:
+    def firmware_version(self) -> str | None:
         """Return the software version for this device."""
-        return self._sw_build_id
-
-    @sw_version.setter
-    def sw_version(self, sw_build_id: int) -> None:
-        """Set the software version for this device."""
-        self._sw_build_id = sw_build_id
+        return self._firmware_version
 
     @property
     def platform_entities(self) -> dict[tuple[Platform, str], PlatformEntity]:
@@ -587,9 +604,21 @@ def new(
         """Create new device."""
         return cls(zigpy_dev, gateway)
 
-    def async_update_sw_build_id(self, sw_version: int) -> None:
-        """Update device sw version."""
-        self._sw_build_id = sw_version
+    def async_update_firmware_version(self, firmware_version: str) -> None:
+        """Update device firmware version."""
+        if firmware_version == self._firmware_version:
+            return
+
+        old_firmware_version = self._firmware_version
+        self._firmware_version = firmware_version
+
+        self.emit(
+            DeviceFirmwareInfoUpdatedEvent.event_type,
+            DeviceFirmwareInfoUpdatedEvent(
+                old_firmware_version=old_firmware_version,
+                new_firmware_version=firmware_version,
+            ),
+        )
 
     async def _check_available(self, *_: Any) -> None:
         # don't flip the availability state of the coordinator
@@ -903,16 +932,32 @@ async def async_initialize(self, from_cache: bool = False) -> None:
         # At this point we can compute a primary entity
         self._compute_primary_entity()
 
+        # Sync the device's firmware version with the first platform entity
+        for (platform, _unique_id), entity in self.platform_entities.items():
+            if platform != Platform.UPDATE:
+                continue
+
+            self._firmware_version = entity.installed_version
+
+            def entity_update_listener(event: EntityStateChangedEvent) -> None:
+                """Listen to firmware update entity changes."""
+                entity = self.get_platform_entity(event.platform, event.unique_id)
+                self.async_update_firmware_version(entity.installed_version)
+
+            self._on_remove_callbacks.append(
+                entity.on_event(STATE_CHANGED, entity_update_listener)
+            )
+
+            break
+
         self.debug("power source: %s", self.power_source)
         self.status = DeviceStatus.INITIALIZED
         self.debug("completed initialization")
 
     async def on_remove(self) -> None:
         """Cancel tasks this device owns."""
-        self._zdo_handler.on_remove()
-
-        for endpoint in self._endpoints.values():
-            endpoint.on_remove()
+        for callback in self._on_remove_callbacks:
+            callback()
 
         for platform_entity in self._platform_entities.values():
             await platform_entity.on_remove()