Properly tear down group entities on membership changes and group removal

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

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

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,

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