From 80b8570d21804f50dd6d96e699a72a5260dfcda2 Mon Sep 17 00:00:00 2001 From: David Mulcahey Date: Tue, 14 Jul 2026 09:42:45 -0400 Subject: [PATCH] Harden device entity lifecycle and command forwarding --- tests/test_device.py | 214 ++++++++++++++++++++++++++ zha/application/platforms/__init__.py | 11 +- zha/zigbee/device.py | 28 +++- 3 files changed, 250 insertions(+), 3 deletions(-) diff --git a/tests/test_device.py b/tests/test_device.py index 8c9703555..531c665e3 100644 --- a/tests/test_device.py +++ b/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 Status, WriteAttributesResponse from zigpy.zcl.helpers import ReportingConfig import zigpy.zdo.types as zdo_t @@ -598,6 +600,93 @@ async def test_issue_cluster_command( assert cluster.request.await_count == 1 +@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, @@ -1546,6 +1635,131 @@ 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) + + 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_async_initialize_drains_stale_pending_entities_before_discovery( + zha_gateway: Gateway, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test stale pending entities are removed before a new discovery cycle.""" + zigpy_dev = zigpy_device(zha_gateway, with_basic_cluster=True) + zha_device = await join_zigpy_device(zha_gateway, zigpy_dev) + + stale_entities = [mock.Mock(), mock.Mock()] + stale_entities[0].on_remove = AsyncMock() + stale_entities[1].on_remove = AsyncMock( + side_effect=RuntimeError("stale pending cleanup failed") + ) + zha_device._pending_entities.extend(stale_entities) + + def assert_pending_entities_drained() -> None: + assert not zha_device._pending_entities + + with patch.object( + zha_device, + "_discover_new_entities", + side_effect=assert_pending_entities_drained, + ) as discover_new_entities: + await zha_device.async_initialize(from_cache=False) + + for entity in stale_entities: + entity.on_remove.assert_awaited_once_with() + discover_new_entities.assert_called_once_with() + assert "Failed to remove stale pending entity" in caplog.text + assert "stale pending cleanup failed" in caplog.text + + +async def test_concurrent_async_initialize_serializes_reconciliation( + zha_gateway: Gateway, +) -> None: + """Test concurrent initialization passes cannot share pending entities.""" + zigpy_dev = zigpy_device(zha_gateway, with_basic_cluster=True) + zha_device = zha_gateway.get_or_create_device(zigpy_dev) + + first_reconciliation_started = asyncio.Event() + finish_first_reconciliation = asyncio.Event() + original_add_pending_entities = zha_device._add_pending_entities + add_pending_call_count = 0 + + async def controlled_add_pending_entities(*, emit_event: bool = True) -> None: + nonlocal add_pending_call_count + add_pending_call_count += 1 + if add_pending_call_count == 1: + first_reconciliation_started.set() + await finish_first_reconciliation.wait() + await original_add_pending_entities(emit_event=emit_event) + + initialize_tasks: list[asyncio.Task[None]] = [] + with ( + patch.object( + zha_device, + "_discover_new_entities", + wraps=zha_device._discover_new_entities, + ) as discover_new_entities, + patch.object( + zha_device, + "_add_pending_entities", + side_effect=controlled_add_pending_entities, + ), + ): + try: + first_initialize = asyncio.create_task( + zha_device.async_initialize(from_cache=False) + ) + initialize_tasks.append(first_initialize) + await asyncio.wait_for(first_reconciliation_started.wait(), timeout=1) + + second_initialize = asyncio.create_task( + zha_device.async_initialize(from_cache=False) + ) + initialize_tasks.append(second_initialize) + await asyncio.sleep(0) + + assert discover_new_entities.call_count == 1 + assert not second_initialize.done() + + finish_first_reconciliation.set() + await asyncio.gather(*initialize_tasks) + + assert discover_new_entities.call_count == 2 + assert not zha_device._pending_entities + finally: + finish_first_reconciliation.set() + for task in initialize_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*initialize_tasks, return_exceptions=True) + + async def test_initial_entity_discovery_does_not_emit_events( zha_gateway: Gateway, ) -> None: diff --git a/zha/application/platforms/__init__.py b/zha/application/platforms/__init__.py index 675918e88..151919d50 100644 --- a/zha/application/platforms/__init__.py +++ b/zha/application/platforms/__init__.py @@ -445,7 +445,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) @@ -455,6 +462,8 @@ async def on_remove(self) -> None: for task in tasks: self.debug("Cancelling task: %s", task) task.cancel() + with suppress(ValueError): + self._tracked_tasks.remove(task) with suppress(asyncio.CancelledError): await asyncio.gather(*tasks, return_exceptions=True) diff --git a/zha/zigbee/device.py b/zha/zigbee/device.py index d9a6a1cb0..3e6f2df57 100644 --- a/zha/zigbee/device.py +++ b/zha/zigbee/device.py @@ -369,6 +369,7 @@ def __init__( self.unique_id = str(zigpy_device.ieee) self._gateway: Gateway = _gateway + self._initialize_lock = asyncio.Lock() self._platform_entities: dict[tuple[Platform, str], PlatformEntity] = {} self._pending_entities: list[PlatformEntity] = [] # All entities discovered for this device, including ones removed by a quirk. @@ -1206,8 +1207,25 @@ async def recompute_entities(self) -> None: async def async_initialize(self, from_cache: bool = False) -> None: """Initialize cluster handlers.""" + async with self._initialize_lock: + await self._async_initialize(from_cache) + + async def _async_initialize(self, from_cache: bool) -> None: + """Initialize cluster handlers while holding the initialization lock.""" self.debug("started initialization") + while self._pending_entities: + entity = self._pending_entities.pop() + try: + await entity.on_remove() + except Exception: + _LOGGER.warning( + "Failed to remove stale pending entity %s for device %s", + entity, + self, + exc_info=True, + ) + # We discover prospective entities before initialization self._discover_new_entities() @@ -1421,6 +1439,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( ( @@ -1430,11 +1451,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",