diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 17792ba40..5f72aa93b 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -41,6 +41,7 @@ from zha.application.helpers import ZHAData from zha.application.platforms import 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 @@ -1285,3 +1286,140 @@ async def test_group_on_remove_entity_failure( assert "Failed to remove group entity" in caplog.text assert "Group entity removal failed" in caplog.text + + +async def test_group_entities_removed_when_group_shrinks( + zha_gateway: Gateway, +) -> None: + """Test group entities are torn down when a group drops below 2 members.""" + device_light_1 = await device_light_1_mock(zha_gateway) + device_light_2 = await device_light_2_mock(zha_gateway) + + device_1_light_entity = get_entity(device_light_1, platform=Platform.LIGHT) + baseline_listener_count = len( + device_1_light_entity._listeners.get(STATE_CHANGED, []) + ) + + zha_group: Group = await zha_gateway.async_create_zigpy_group( + "Test Group", + [ + GroupMemberReference(ieee=device_light_1.ieee, endpoint_id=1), + GroupMemberReference(ieee=device_light_2.ieee, endpoint_id=1), + ], + ) + await zha_gateway.async_block_till_done() + + group_entity = get_group_entity(zha_group, platform=Platform.LIGHT) + assert group_entity is not None + assert zha_group._entity_unsubs + # the group entity is subscribed to member entity state changes + assert ( + len(device_1_light_entity._listeners[STATE_CHANGED]) + == baseline_listener_count + 1 + ) + + await zha_group.async_remove_members( + [GroupMemberReference(ieee=device_light_2.ieee, endpoint_id=1)] + ) + await zha_gateway.async_block_till_done() + + assert len(zha_group.members) == 1 + assert not zha_group.group_entities + # all entity subscriptions are released + assert not zha_group._entity_unsubs + assert ( + len(device_1_light_entity._listeners[STATE_CHANGED]) == baseline_listener_count + ) + + +async def test_group_platform_entity_removed_when_platform_drops_below_two( + zha_gateway: Gateway, +) -> None: + """Test a platform's group entity is removed when it loses eligibility.""" + device_light_1 = await device_light_1_mock(zha_gateway) + device_light_2 = await device_light_2_mock(zha_gateway) + + switch_endpoint_config = { + 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 = await join_zigpy_device( + zha_gateway, + create_mock_zigpy_device( + zha_gateway, switch_endpoint_config, ieee="00:0d:6f:00:0a:90:69:e1" + ), + ) + device_switch_2 = await join_zigpy_device( + zha_gateway, + create_mock_zigpy_device( + zha_gateway, switch_endpoint_config, ieee="00:0d:6f:00:0a:90:69:e2" + ), + ) + + zha_group: Group = await zha_gateway.async_create_zigpy_group( + "Test Group", + [ + 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), + ], + ) + await zha_gateway.async_block_till_done() + + light_group_entity = get_group_entity(zha_group, platform=Platform.LIGHT) + switch_group_entity = get_group_entity(zha_group, platform=Platform.SWITCH) + assert light_group_entity is not None + assert switch_group_entity is not None + + # drop the switch platform below 2 eligible members + await zha_group.async_remove_members( + [GroupMemberReference(ieee=device_switch_2.ieee, endpoint_id=1)] + ) + await zha_gateway.async_block_till_done() + + assert len(zha_group.members) == 3 + with pytest.raises(KeyError): + get_group_entity(zha_group, platform=Platform.SWITCH) + # the light group entity is preserved + assert get_group_entity(zha_group, platform=Platform.LIGHT) is light_group_entity + + +async def test_group_removed_tears_down_group_entities( + zha_gateway: Gateway, +) -> None: + """Test group entities are torn down when the zigpy group is removed.""" + device_light_1 = await device_light_1_mock(zha_gateway) + device_light_2 = await device_light_2_mock(zha_gateway) + + device_1_light_entity = get_entity(device_light_1, platform=Platform.LIGHT) + baseline_listener_count = len( + device_1_light_entity._listeners.get(STATE_CHANGED, []) + ) + + zha_group: Group = await zha_gateway.async_create_zigpy_group( + "Test Group", + [ + GroupMemberReference(ieee=device_light_1.ieee, endpoint_id=1), + GroupMemberReference(ieee=device_light_2.ieee, endpoint_id=1), + ], + ) + await zha_gateway.async_block_till_done() + + assert get_group_entity(zha_group, platform=Platform.LIGHT) is not None + + # remove the zigpy group directly, without removing the members first + zha_gateway.application_controller.groups.pop(zha_group.group_id) + await zha_gateway.async_block_till_done() + + assert zha_gateway.get_group(zha_group.group_id) is None + assert not zha_group.group_entities + # all entity subscriptions are released + assert not zha_group._entity_unsubs + assert ( + len(device_1_light_entity._listeners[STATE_CHANGED]) == baseline_listener_count + ) diff --git a/zha/application/discovery.py b/zha/application/discovery.py index c5ce8bc8f..a471804c2 100644 --- a/zha/application/discovery.py +++ b/zha/application/discovery.py @@ -38,6 +38,7 @@ update, virtual, ) +from zha.async_ import create_eager_task from zha.zigbee.group import Group if TYPE_CHECKING: @@ -107,6 +108,21 @@ def inner(*args: P.args, **kwargs: P.kwargs) -> Iterator[T]: return inner +def _remove_stale_group_entities( + group: Group, stale_entities: list[GroupEntity] +) -> None: + """Unregister stale group entities and schedule their removal.""" + for entity in stale_entities: + _LOGGER.info("Removing stale group entity %s from group %s", entity, group.name) + group.unregister_group_entity(entity) + group.gateway.track_task( + create_eager_task( + entity.on_remove(), name=f"remove_group_entity_{entity.unique_id}" + ) + ) + group.update_entity_subscriptions() + + @ignore_exceptions_during_iteration def discover_group_entities(group: Group) -> Iterator[GroupEntity]: """Process a group and create any entities that are needed.""" @@ -117,7 +133,8 @@ def discover_group_entities(group: Group) -> Iterator[GroupEntity]: group.name, group.group_id, ) - group.group_entities.clear() + if group.group_entities: + _remove_stale_group_entities(group, list(group.group_entities.values())) return # We only create groups with two or more devices @@ -130,10 +147,19 @@ def discover_group_entities(group: Group) -> Iterator[GroupEntity]: for entity in member.associated_entities: platform_counts[entity.PLATFORM] += 1 - for platform, count in platform_counts.items(): - if count < 2: - continue + eligible_platforms = { + platform for platform, count in platform_counts.items() if count >= 2 + } + + # Remove group entities for platforms that no longer have enough members + if stale_entities := [ + entity + for entity in group.group_entities.values() + if entity.PLATFORM not in eligible_platforms + ]: + _remove_stale_group_entities(group, stale_entities) + for platform in eligible_platforms: for group_entity_class in GROUP_ENTITY_REGISTRY: if platform != group_entity_class.PLATFORM: continue diff --git a/zha/application/gateway.py b/zha/application/gateway.py index ecff916d1..03437bd62 100644 --- a/zha/application/gateway.py +++ b/zha/application/gateway.py @@ -648,6 +648,9 @@ def group_removed(self, zigpy_group: zigpy.group.Group) -> None: self._emit_group_gateway_message(zigpy_group, ZHA_GW_MSG_GROUP_REMOVED) zha_group = self._groups.pop(zigpy_group.group_id) zha_group.info("group_removed") + self.track_task( + create_eager_task(zha_group.on_remove(), name="Gateway.group_removed") + ) def _emit_group_gateway_message( # pylint: disable=unused-argument self, diff --git a/zha/zigbee/group.py b/zha/zigbee/group.py index 9f16b40af..34034ce7a 100644 --- a/zha/zigbee/group.py +++ b/zha/zigbee/group.py @@ -232,7 +232,9 @@ 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: + # Only unregister if this exact entity is registered, so a delayed + # removal cannot unregister a recreated entity with the same unique id + if self._group_entities.get(group_entity.unique_id) is group_entity: self._group_entities.pop(group_entity.unique_id) self._entity_unsubs.pop(group_entity.unique_id)() @@ -354,3 +356,7 @@ async def on_remove(self) -> None: group_entity, exc_info=True, ) + # Unsubscribe any remaining member entity subscriptions + while self._entity_unsubs: + _, unsub = self._entity_unsubs.popitem() + unsub()