diff --git a/tests/test_device.py b/tests/test_device.py index 8c9703555..d8d80061a 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -124,6 +124,28 @@ def zigpy_device_mains(zha_gateway: Gateway, with_basic_cluster: bool = True): ) +def test_discover_entities_continues_after_endpoint_exception() -> None: + """Test discovery preserves partial results and continues after an exception.""" + zha_device = mock.Mock(ieee="00:0d:6f:00:0a:90:69:e7", is_coordinator=False) + endpoint_1 = mock.Mock(id=1, device=zha_device) + endpoint_2 = mock.Mock(id=2, device=zha_device) + zha_device.endpoints = {1: endpoint_1, 2: endpoint_2} + + def discover_endpoint(endpoint): + if endpoint.id == 1: + yield mock.sentinel.partial_entity + raise RuntimeError("endpoint discovery failed") + yield mock.sentinel.later_entity + + with mock.patch( + "zha.application.discovery.discover_entities_for_endpoint", + side_effect=discover_endpoint, + ): + entities = list(Device.discover_entities(zha_device)) + + assert entities == [mock.sentinel.partial_entity, mock.sentinel.later_entity] + + async def _send_time_changed(zha_gateway: Gateway, seconds: int): """Send a time changed event.""" await asyncio.sleep(seconds) diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 17792ba40..4574403a7 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -1,6 +1,7 @@ """Test ZHA Gateway.""" import asyncio +from collections.abc import Callable from contextlib import suppress from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch @@ -39,8 +40,9 @@ RawDeviceInitializedEvent, ) from zha.application.helpers import ZHAData -from zha.application.platforms import GroupEntity +from zha.application.platforms import BaseEntity, GroupEntity from zha.application.platforms.light.const import EFFECT_OFF, LightEntityFeature +from zha.const import STATE_CHANGED from zha.quirks import DeviceMatch, DeviceRegistry, ModelInfo, QuirkRegistryEntry from zha.zigbee.device import Device, DeviceEntityAddedEvent, DeviceEntityRemovedEvent from zha.zigbee.group import Group, GroupMemberReference @@ -151,6 +153,194 @@ async def device_light_2_mock(zha_gateway: Gateway) -> Device: return zha_device +def _has_state_changed_listener( + entity: BaseEntity, callback: Callable[..., object] +) -> bool: + """Return whether an entity has the given state listener.""" + return any( + listener.callback == callback + for listener in entity._listeners.get(STATE_CHANGED, []) + ) + + +async def _create_light_group( + zha_gateway: Gateway, +) -> tuple[Group, Device, Device, GroupEntity]: + """Create a two-member light group.""" + device_light_1 = await device_light_1_mock(zha_gateway) + device_light_2 = await device_light_2_mock(zha_gateway) + members = [ + GroupMemberReference(ieee=device_light_1.ieee, endpoint_id=1), + GroupMemberReference(ieee=device_light_2.ieee, endpoint_id=1), + ] + zha_group = await zha_gateway.async_create_zigpy_group("Test Group", members) + assert zha_group is not None + await zha_gateway.async_block_till_done() + entity = get_group_entity(zha_group, platform=Platform.LIGHT) + return zha_group, device_light_1, device_light_2, entity + + +async def test_direct_group_membership_change_refreshes_entity_subscriptions( + zha_gateway: Gateway, +) -> None: + """Test direct membership changes refresh an existing group entity.""" + zha_group, _, _, entity = await _create_light_group(zha_gateway) + device_light_3 = await join_zigpy_device( + zha_gateway, + create_mock_zigpy_device( + zha_gateway, + { + 1: { + SIG_EP_INPUT: [ + general.OnOff.cluster_id, + general.LevelControl.cluster_id, + lighting.Color.cluster_id, + general.Groups.cluster_id, + ], + SIG_EP_OUTPUT: [], + SIG_EP_TYPE: zha.DeviceType.COLOR_DIMMABLE_LIGHT, + SIG_EP_PROFILE: zha.PROFILE_ID, + } + }, + ieee="03:2d:6f:00:0a:90:69:e8", + ), + ) + third_member_entity = get_entity(device_light_3, platform=Platform.LIGHT) + endpoint = device_light_3.device.endpoints[1] + + assert not _has_state_changed_listener(third_member_entity, entity.debounced_update) + + zha_group.zigpy_group.add_member(endpoint) + + assert zha_group.group_entities[entity.unique_id] is entity + assert _has_state_changed_listener(third_member_entity, entity.debounced_update) + + zha_group.zigpy_group.remove_member(endpoint) + + assert zha_group.group_entities[entity.unique_id] is entity + assert not _has_state_changed_listener(third_member_entity, entity.debounced_update) + + +async def test_delayed_group_entity_cleanup_cannot_remove_replacement( + zha_gateway: Gateway, +) -> None: + """Test delayed stale cleanup cannot unregister a rediscovered entity.""" + zha_group, device_light_1, device_light_2, old_entity = await _create_light_group( + zha_gateway + ) + member_entities = [ + get_entity(device, platform=Platform.LIGHT) + for device in (device_light_1, device_light_2) + ] + original_on_remove = old_entity.on_remove + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + + async def delayed_cleanup() -> None: + cleanup_started.set() + await release_cleanup.wait() + await original_on_remove() + + try: + with patch.object( + old_entity, + "on_remove", + new=AsyncMock(side_effect=delayed_cleanup), + ) as on_remove: + endpoint = device_light_2.device.endpoints[1] + zha_group.zigpy_group.remove_member(endpoint) + await cleanup_started.wait() + + assert all( + not _has_state_changed_listener(member, old_entity.debounced_update) + for member in member_entities + ) + + zha_group.zigpy_group.add_member(endpoint) + replacement = get_group_entity(zha_group, platform=Platform.LIGHT) + + assert replacement is not old_entity + assert [ + entity + for entity in zha_group.group_entities.values() + if entity.PLATFORM == Platform.LIGHT + ] == [replacement] + assert all( + _has_state_changed_listener(member, replacement.debounced_update) + for member in member_entities + ) + + release_cleanup.set() + await zha_gateway.async_block_till_done() + + assert zha_group.group_entities[replacement.unique_id] is replacement + assert on_remove.await_count == 1 + assert all( + _has_state_changed_listener(member, replacement.debounced_update) + for member in member_entities + ) + finally: + release_cleanup.set() + await zha_gateway.async_block_till_done() + await zha_group.on_remove() + + +async def test_group_platform_quorum_loss_prunes_only_stale_entity( + zha_gateway: Gateway, +) -> None: + """Test losing one platform quorum prunes only that group entity.""" + device_light_1 = await device_light_1_mock(zha_gateway) + device_light_2 = await device_light_2_mock(zha_gateway) + switch_endpoints = { + 1: { + SIG_EP_INPUT: [general.OnOff.cluster_id, general.Groups.cluster_id], + SIG_EP_OUTPUT: [], + SIG_EP_TYPE: zha.DeviceType.ON_OFF_SWITCH, + SIG_EP_PROFILE: zha.PROFILE_ID, + } + } + device_switch_1, device_switch_2 = [ + await join_zigpy_device( + zha_gateway, + create_mock_zigpy_device(zha_gateway, switch_endpoints, ieee=ieee), + ) + for ieee in ("03:2d:6f:00:0a:90:69:e8", "04:2d:6f:00:0a:90:69:e8") + ] + members = [ + GroupMemberReference(ieee=device_light_1.ieee, endpoint_id=1), + GroupMemberReference(ieee=device_light_2.ieee, endpoint_id=1), + GroupMemberReference(ieee=device_switch_1.ieee, endpoint_id=1), + GroupMemberReference(ieee=device_switch_2.ieee, endpoint_id=1), + ] + zha_group = await zha_gateway.async_create_zigpy_group("Test Group", members) + await zha_gateway.async_block_till_done() + light_entity = get_group_entity(zha_group, platform=Platform.LIGHT) + switch_entity = get_group_entity(zha_group, platform=Platform.SWITCH) + light_member_entity = get_entity(device_light_1, platform=Platform.LIGHT) + switch_member_entity = get_entity(device_switch_1, platform=Platform.SWITCH) + + try: + with patch.object( + light_entity, + "on_remove", + new=AsyncMock(wraps=light_entity.on_remove), + ) as on_remove: + await zha_group.async_remove_members([members[1]]) + await zha_gateway.async_block_till_done() + + assert on_remove.await_count == 1 + assert light_entity.unique_id not in zha_group.group_entities + assert zha_group.group_entities[switch_entity.unique_id] is switch_entity + assert not _has_state_changed_listener( + light_member_entity, light_entity.debounced_update + ) + assert _has_state_changed_listener( + switch_member_entity, switch_entity.debounced_update + ) + finally: + await zha_group.on_remove() + + async def test_device_left(zha_gateway: Gateway) -> None: """Device leaving the network should become unavailable.""" zigpy_dev_basic = create_mock_zigpy_device(zha_gateway, ZIGPY_DEVICE_BASIC) @@ -1277,11 +1467,31 @@ async def test_group_on_remove_entity_failure( await zha_gateway.async_block_till_done() group_entity = get_group_entity(zha_group, platform=Platform.LIGHT) + member_entity = get_entity(zha_device_1, platform=Platform.LIGHT) + original_on_remove = group_entity.on_remove - with patch.object( - group_entity, "on_remove", side_effect=Exception("Group entity removal failed") - ): - await zha_group.on_remove() - - assert "Failed to remove group entity" in caplog.text - assert "Group entity removal failed" in caplog.text + try: + assert _has_state_changed_listener( + group_entity, zha_group._handle_maybe_update_group_members + ) + assert _has_state_changed_listener(member_entity, group_entity.debounced_update) + + with patch.object( + group_entity, + "on_remove", + new=AsyncMock(side_effect=Exception("Group entity removal failed")), + ) as on_remove: + await zha_group.on_remove() + + assert on_remove.await_count == 1 + assert group_entity.unique_id not in zha_group.group_entities + assert not _has_state_changed_listener( + group_entity, zha_group._handle_maybe_update_group_members + ) + assert not _has_state_changed_listener( + member_entity, group_entity.debounced_update + ) + assert "Failed to remove group entity" in caplog.text + assert "Group entity removal failed" in caplog.text + finally: + await original_on_remove() diff --git a/zha/application/discovery.py b/zha/application/discovery.py index c5ce8bc8f..95a1e27c3 100644 --- a/zha/application/discovery.py +++ b/zha/application/discovery.py @@ -117,7 +117,8 @@ def discover_group_entities(group: Group) -> Iterator[GroupEntity]: group.name, group.group_id, ) - group.group_entities.clear() + for group_entity in tuple(group.group_entities.values()): + group.schedule_group_entity_cleanup(group_entity) return # We only create groups with two or more devices @@ -130,8 +131,16 @@ def discover_group_entities(group: Group) -> Iterator[GroupEntity]: for entity in member.associated_entities: platform_counts[entity.PLATFORM] += 1 + for group_entity in tuple(group.group_entities.values()): + if platform_counts[group_entity.PLATFORM] < 2: + group.schedule_group_entity_cleanup(group_entity) + + existing_platforms = { + group_entity.PLATFORM for group_entity in group.group_entities.values() + } + for platform, count in platform_counts.items(): - if count < 2: + if count < 2 or platform in existing_platforms: continue for group_entity_class in GROUP_ENTITY_REGISTRY: @@ -142,6 +151,7 @@ def discover_group_entities(group: Group) -> Iterator[GroupEntity]: group_entity_class, group.name, ) + existing_platforms.add(platform) yield group_entity_class(group) diff --git a/zha/application/gateway.py b/zha/application/gateway.py index ecff916d1..096ac7845 100644 --- a/zha/application/gateway.py +++ b/zha/application/gateway.py @@ -618,6 +618,7 @@ def group_member_removed( for entity in discovery.discover_group_entities(zha_group): entity.on_add() + zha_group.update_entity_subscriptions() zha_group.info("group_member_removed - endpoint: %s", endpoint) self._emit_group_gateway_message(zigpy_group, ZHA_GW_MSG_GROUP_MEMBER_REMOVED) @@ -632,6 +633,7 @@ def group_member_added( for entity in discovery.discover_group_entities(zha_group): entity.on_add() + zha_group.update_entity_subscriptions() zha_group.info("group_member_added - endpoint: %s", endpoint) self._emit_group_gateway_message(zigpy_group, ZHA_GW_MSG_GROUP_MEMBER_ADDED) diff --git a/zha/zigbee/device.py b/zha/zigbee/device.py index d9a6a1cb0..7dfbd53e6 100644 --- a/zha/zigbee/device.py +++ b/zha/zigbee/device.py @@ -1070,7 +1070,14 @@ def discover_entities(self) -> Iterator[BaseEntity]: str(endpoint.device.ieee), endpoint.id, ) - yield from discovery.discover_entities_for_endpoint(endpoint) + try: + yield from discovery.discover_entities_for_endpoint(endpoint) + except Exception: # pylint: disable=broad-except + _LOGGER.exception( + "Failed to discover entities for endpoint %s on device %s", + endpoint.id, + str(self.ieee), + ) def _discover_new_entities(self) -> None: self._discovered_entities.clear() diff --git a/zha/zigbee/group.py b/zha/zigbee/group.py index 9f16b40af..f5f6dba61 100644 --- a/zha/zigbee/group.py +++ b/zha/zigbee/group.py @@ -157,6 +157,7 @@ def __init__( self._zigpy_group = zigpy_group self._group_entities: dict[str, GroupEntity] = {} self._entity_unsubs: dict[str, Callable] = {} + self._group_entity_cleanup_tasks: set[asyncio.Task[None]] = set() @property def name(self) -> str: @@ -232,9 +233,35 @@ def register_group_entity(self, group_entity: GroupEntity) -> None: def unregister_group_entity(self, group_entity: GroupEntity) -> None: """Unregister a group entity.""" - if group_entity.unique_id in self._group_entities: - self._group_entities.pop(group_entity.unique_id) - self._entity_unsubs.pop(group_entity.unique_id)() + if self._group_entities.get(group_entity.unique_id) is not group_entity: + return + + self._group_entities.pop(group_entity.unique_id) + if unsubscribe := self._entity_unsubs.pop(group_entity.unique_id, None): + unsubscribe() + self.update_entity_subscriptions() + + async def _async_cleanup_group_entity(self, group_entity: GroupEntity) -> None: + """Run and observe lifecycle cleanup for a group entity.""" + try: + await group_entity.on_remove() + except Exception: + _LOGGER.warning( + "Failed to remove group entity %s", + group_entity, + exc_info=True, + ) + + def schedule_group_entity_cleanup(self, group_entity: GroupEntity) -> None: + """Detach a group entity and schedule its lifecycle cleanup.""" + self.unregister_group_entity(group_entity) + cleanup_task = self.gateway.async_create_task( + self._async_cleanup_group_entity(group_entity), + name=f"zha.group-remove-entity-{group_entity.unique_id}", + eager_start=True, + ) + self._group_entity_cleanup_tasks.add(cleanup_task) + cleanup_task.add_done_callback(self._group_entity_cleanup_tasks.discard) def _handle_maybe_update_group_members(self, event: EntityStateChangedEvent): """Handle the maybe update group members event.""" @@ -346,11 +373,7 @@ def log(self, level: int, msg: str, *args: Any, **kwargs) -> None: async def on_remove(self) -> None: """Cancel tasks this group owns.""" for group_entity in tuple(self._group_entities.values()): - try: - await group_entity.on_remove() - except Exception: - _LOGGER.warning( - "Failed to remove group entity %s", - group_entity, - exc_info=True, - ) + self.schedule_group_entity_cleanup(group_entity) + + if cleanup_tasks := tuple(self._group_entity_cleanup_tasks): + await asyncio.gather(*cleanup_tasks)