Preserve entity cleanup and explicit command manufacturers

tests/test_device.py

@@ -3,6 +3,7 @@
 import asyncio
 import logging
 import time
+from typing import Any
 from unittest import mock
 from unittest.mock import AsyncMock, call, patch
 
@@ -22,6 +23,7 @@
 from zigpy.zcl.clusters.general import Ota, PowerConfiguration
 from zigpy.zcl.clusters.lighting import Color
 from zigpy.zcl.clusters.measurement import CarbonDioxideConcentration
+from zigpy.zcl.clusters.wwah import WorksWithAllHubs
 from zigpy.zcl.foundation import (
     GENERAL_COMMANDS,
     GeneralCommand,
@@ -770,6 +772,93 @@ def _response(status: lightlink.Status) -> network_start_rsp:
         )
 
 
+@pytest.mark.parametrize("use_args", [True, False], ids=["args", "params"])
+@pytest.mark.parametrize(
+    "command_type",
+    [CLUSTER_COMMAND_SERVER, CLUSTER_COMMANDS_CLIENT],
+    ids=["server", "client"],
+)
+async def test_issue_cluster_command_forwards_manufacturer(
+    zha_gateway: Gateway,
+    command_type: str,
+    use_args: bool,
+) -> None:
+    """Test manufacturer forwarding for every cluster command invocation path."""
+    zigpy_dev = zigpy_device(zha_gateway, with_basic_cluster=True)
+    zigpy_dev.endpoints[3].add_input_cluster(general.Groups.cluster_id)
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    if command_type == CLUSTER_COMMAND_SERVER:
+        command = general.Groups.ServerCommandDefs.add.id
+        command_args: list[Any] = [0x0001, "test group"]
+        command_params: dict[str, Any] = {
+            "group_id": 0x0001,
+            "group_name": "test group",
+        }
+        transport_patch = "zigpy.zcl.Cluster.request"
+        transport_response = [0x05, Status.SUCCESS]
+    else:
+        command = general.Groups.ClientCommandDefs.add_response.id
+        command_args = [Status.SUCCESS, 0x0001]
+        command_params = {"status": Status.SUCCESS, "group_id": 0x0001}
+        transport_patch = "zigpy.zcl.Cluster.reply"
+        transport_response = None
+
+    args = command_args if use_args else None
+    params = None if use_args else command_params
+
+    with patch(transport_patch, return_value=transport_response) as transport:
+        await zha_device.issue_cluster_command(
+            3,
+            general.Groups.cluster_id,
+            command,
+            command_type,
+            args,
+            params,
+            manufacturer=0,
+        )
+
+        assert transport.await_count == 1
+        assert transport.await_args.kwargs["manufacturer"] == 0
+
+
+@pytest.mark.parametrize("use_args", [True, False], ids=["args", "params"])
+async def test_issue_cluster_command_preserves_manufacturer_inference(
+    zha_gateway: Gateway,
+    use_args: bool,
+) -> None:
+    """Test manufacturer inference when no manufacturer is passed."""
+    zigpy_dev = zigpy_device(zha_gateway, with_basic_cluster=True)
+    zigpy_dev.endpoints[3].add_input_cluster(WorksWithAllHubs.cluster_id)
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    command_def = (
+        WorksWithAllHubs.ClientCommandDefs.aps_link_key_authorization_query_response
+    )
+    command_args: list[Any] = [general.OnOff.cluster_id, True]
+    command_params: dict[str, Any] = {
+        "cluster_id": general.OnOff.cluster_id,
+        "aps_link_key_auth_status": True,
+    }
+    args = command_args if use_args else None
+    params = None if use_args else command_params
+
+    with patch("zigpy.zcl.Cluster.reply", return_value=None) as transport:
+        await zha_device.issue_cluster_command(
+            3,
+            WorksWithAllHubs.cluster_id,
+            command_def.id,
+            CLUSTER_COMMANDS_CLIENT,
+            args,
+            params,
+        )
+
+        assert transport.await_count == 1
+        assert (
+            transport.await_args.kwargs["manufacturer"] == command_def.manufacturer_code
+        )
+
+
 async def test_async_add_to_group_remove_from_group(
     zha_gateway: Gateway,
     caplog: pytest.LogCaptureFixture,
@@ -1872,6 +1961,38 @@ async def test_device_on_remove_pending_entity_failure(
     assert "Pending entity removal failed" in caplog.text
 
 
+async def test_platform_entity_on_remove_callback_failure_does_not_abort_cleanup(
+    zha_gateway: Gateway,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that a failed remove callback does not abort remaining cleanup."""
+    zigpy_dev = zigpy_device(zha_gateway, with_basic_cluster=True)
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+    entity = get_entity(zha_device, platform=Platform.SWITCH)
+
+    successful_callback = mock.Mock()
+    failing_callback = mock.Mock(side_effect=RuntimeError("entity callback failure"))
+    entity._on_remove_callbacks.extend((successful_callback, failing_callback))
+
+    tracked_task = asyncio.create_task(asyncio.Event().wait())
+    entity._tracked_tasks.append(tracked_task)
+    tracked_task.add_done_callback(entity._tracked_tasks.remove)
+
+    try:
+        await entity.on_remove()
+
+        failing_callback.assert_called_once_with()
+        successful_callback.assert_called_once_with()
+        assert tracked_task.cancelled()
+        assert tracked_task not in entity._tracked_tasks
+        assert "Failed to execute on_remove callback" in caplog.text
+        assert "entity callback failure" in caplog.text
+    finally:
+        if not tracked_task.done():
+            tracked_task.cancel()
+        await asyncio.gather(tracked_task, return_exceptions=True)
+
+
 async def test_initial_entity_discovery_does_not_emit_events(
     zha_gateway: Gateway,
 ) -> None:

zha/application/platforms/__init__.py

@@ -462,7 +462,14 @@ async def on_remove(self) -> None:
         while self._on_remove_callbacks:
             callback = self._on_remove_callbacks.pop()
             self.debug("Running remove callback: %s", callback)
-            callback()
+            try:
+                callback()
+            except Exception:
+                self.warning(
+                    "Failed to execute on_remove callback %s",
+                    callback,
+                    exc_info=True,
+                )
 
         for handle in self._tracked_handles:
             self.debug("Cancelling handle: %s", handle)

zha/zigbee/device.py

@@ -1444,6 +1444,9 @@ async def issue_cluster_command(
             if command_type == CLUSTER_COMMAND_SERVER
             else cluster.client_commands
         )
+        manufacturer_kwargs = (
+            {} if manufacturer is None else {"manufacturer": manufacturer}
+        )
         if args is not None:
             self.warning(
                 (
@@ -1453,11 +1456,14 @@ async def issue_cluster_command(
                 args,
                 [field.name for field in commands[command].schema.fields],
             )
-            response = await getattr(cluster, commands[command].name)(*args)
+            response = await getattr(cluster, commands[command].name)(
+                *args, **manufacturer_kwargs
+            )
         else:
             assert params is not None
             response = await getattr(cluster, commands[command].name)(
-                **convert_to_zcl_values(params, commands[command].schema)
+                **convert_to_zcl_values(params, commands[command].schema),
+                **manufacturer_kwargs,
             )
         self.debug(
             "Issued cluster command: %s %s %s %s %s %s %s %s",