Update sw_version in device registry for ZHA devices (#33184)

* Update firmware version in device registry.
Parse recevied OTA requests for firmware version and update device
registry.

* Update tests.
* Cleanup sw_id_update listener.
* Update ZHA test devices list.

homeassistant/components/zha/core/channels/general.py

@@ -19,6 +19,7 @@ from ..const import (
     SIGNAL_MOVE_LEVEL,
     SIGNAL_SET_LEVEL,
     SIGNAL_STATE_ATTR,
+    SIGNAL_UPDATE_DEVICE,
 )
 from .base import ClientChannel, ZigbeeChannel, parse_and_log_command
 
@@ -333,11 +334,20 @@ class OnOffConfiguration(ZigbeeChannel):
     pass
 
 
+@registries.CLIENT_CHANNELS_REGISTRY.register(general.Ota.cluster_id)
 @registries.ZIGBEE_CHANNEL_REGISTRY.register(general.Ota.cluster_id)
 class Ota(ZigbeeChannel):
     """OTA Channel."""
 
-    pass
+    @callback
+    def cluster_command(
+        self, tsn: int, command_id: int, args: Optional[List[Any]]
+    ) -> None:
+        """Handle OTA commands."""
+        cmd_name = self.cluster.server_commands.get(command_id, [command_id])[0]
+        signal_id = self._ch_pool.unique_id.split("-")[0]
+        if cmd_name == "query_next_image":
+            self.async_send_signal(SIGNAL_UPDATE_DEVICE.format(signal_id), args[3])
 
 
 @registries.ZIGBEE_CHANNEL_REGISTRY.register(general.Partition.cluster_id)

homeassistant/components/zha/core/const.py

@@ -206,6 +206,7 @@ SIGNAL_MOVE_LEVEL = "move_level"
 SIGNAL_REMOVE = "remove"
 SIGNAL_SET_LEVEL = "set_level"
 SIGNAL_STATE_ATTR = "update_state_attribute"
+SIGNAL_UPDATE_DEVICE = "{}_zha_update_device"
 
 UNKNOWN = "unknown"
 UNKNOWN_MANUFACTURER = "unk_manufacturer"

homeassistant/components/zha/core/device.py

@@ -54,6 +54,7 @@ from .const import (
     POWER_BATTERY_OR_UNKNOWN,
     POWER_MAINS_POWERED,
     SIGNAL_AVAILABLE,
+    SIGNAL_UPDATE_DEVICE,
     UNKNOWN,
     UNKNOWN_MANUFACTURER,
     UNKNOWN_MODEL,
@@ -92,8 +93,11 @@ class ZHADevice(LogMixin):
             self.name, self.ieee, SIGNAL_AVAILABLE
         )
         self._checkins_missed_count = 0
-        self._unsub = async_dispatcher_connect(
-            self.hass, self._available_signal, self.async_initialize
+        self.unsubs = []
+        self.unsubs.append(
+            async_dispatcher_connect(
+                self.hass, self._available_signal, self.async_initialize
+            )
         )
         self.quirk_applied = isinstance(self._zigpy_device, zigpy.quirks.CustomDevice)
         self.quirk_class = "{}.{}".format(
@@ -105,8 +109,10 @@ class ZHADevice(LogMixin):
         else:
             self._consider_unavailable_time = _CONSIDER_UNAVAILABLE_BATTERY
         keep_alive_interval = random.randint(*_UPDATE_ALIVE_INTERVAL)
-        self._cancel_available_check = async_track_time_interval(
-            self.hass, self._check_available, timedelta(seconds=keep_alive_interval)
+        self.unsubs.append(
+            async_track_time_interval(
+                self.hass, self._check_available, timedelta(seconds=keep_alive_interval)
+            )
         )
         self._ha_device_id = None
         self.status = DeviceStatus.CREATED
@@ -276,8 +282,24 @@ class ZHADevice(LogMixin):
         """Create new device."""
         zha_dev = cls(hass, zigpy_dev, gateway)
         zha_dev.channels = channels.Channels.new(zha_dev)
+        zha_dev.unsubs.append(
+            async_dispatcher_connect(
+                hass,
+                SIGNAL_UPDATE_DEVICE.format(zha_dev.channels.unique_id),
+                zha_dev.async_update_sw_build_id,
+            )
+        )
         return zha_dev
 
+    @callback
+    def async_update_sw_build_id(self, sw_version: int):
+        """Update device sw version."""
+        if self.device_id is None:
+            return
+        self._zha_gateway.ha_device_registry.async_update_device(
+            self.device_id, sw_version=f"0x{sw_version:08x}"
+        )
+
     async def _check_available(self, *_):
         if self.last_seen is None:
             self.update_available(False)
@@ -370,8 +392,8 @@ class ZHADevice(LogMixin):
     @callback
     def async_cleanup_handles(self) -> None:
         """Unsubscribe the dispatchers and timers."""
-        self._unsub()
-        self._cancel_available_check()
+        for unsubscribe in self.unsubs:
+            unsubscribe()
 
     @callback
     def async_update_last_seen(self, last_seen):

tests/components/zha/test_device.py

@@ -8,9 +8,10 @@ import pytest
 import zigpy.zcl.clusters.general as general
 
 import homeassistant.components.zha.core.device as zha_core_device
+import homeassistant.helpers.device_registry as ha_dev_reg
 import homeassistant.util.dt as dt_util
 
-from .common import async_enable_traffic
+from .common import async_enable_traffic, make_zcl_header
 
 from tests.common import async_fire_time_changed
 
@@ -63,6 +64,26 @@ def device_without_basic_channel(zigpy_device):
     return zigpy_device(with_basic_channel=False)
 
 
+@pytest.fixture
+async def ota_zha_device(zha_device_restored, zigpy_device_mock):
+    """ZHA device with OTA cluster fixture."""
+    zigpy_dev = zigpy_device_mock(
+        {
+            1: {
+                "in_clusters": [general.Basic.cluster_id],
+                "out_clusters": [general.Ota.cluster_id],
+                "device_type": 0x1234,
+            }
+        },
+        "00:11:22:33:44:55:66:77",
+        "test manufacturer",
+        "test model",
+    )
+
+    zha_device = await zha_device_restored(zigpy_dev)
+    return zha_device
+
+
 def _send_time_changed(hass, seconds):
     """Send a time changed event."""
     now = dt_util.utcnow() + timedelta(seconds=seconds)
@@ -190,3 +211,20 @@ async def test_check_available_no_basic_channel(
     await hass.async_block_till_done()
     assert zha_device.available is False
     assert "does not have a mandatory basic cluster" in caplog.text
+
+
+async def test_ota_sw_version(hass, ota_zha_device):
+    """Test device entry gets sw_version updated via OTA channel."""
+
+    ota_ch = ota_zha_device.channels.pools[0].client_channels["1:0x0019"]
+    dev_registry = await ha_dev_reg.async_get_registry(hass)
+    entry = dev_registry.async_get(ota_zha_device.device_id)
+    assert entry.sw_version is None
+
+    cluster = ota_ch.cluster
+    hdr = make_zcl_header(1, global_command=False)
+    sw_version = 0x2345
+    cluster.handle_message(hdr, [1, 2, 3, sw_version, None])
+    await hass.async_block_till_done()
+    entry = dev_registry.async_get(ota_zha_device.device_id)
+    assert int(entry.sw_version, base=16) == sw_version