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):