Cleanup ZHA group entity lifecycle

homeassistant/components/zha/__init__.py

@@ -32,6 +32,7 @@
     SIGNAL_ADD_ENTITIES,
     RadioType,
 )
+from .core.discovery import GROUP_PROBE
 
 DEVICE_CONFIG_SCHEMA_ENTRY = vol.Schema({vol.Optional(ha_const.CONF_TYPE): cv.string})
 
@@ -138,6 +139,7 @@ async def async_unload_entry(hass, config_entry):
     """Unload ZHA config entry."""
     await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].shutdown()
 
+    GROUP_PROBE.cleanup()
     api.async_unload_api(hass)
 
     dispatchers = hass.data[DATA_ZHA].get(DATA_ZHA_DISPATCHERS, [])

homeassistant/components/zha/core/const.py

@@ -213,6 +213,7 @@ def list(cls):
 SIGNAL_STATE_ATTR = "update_state_attribute"
 SIGNAL_UPDATE_DEVICE = "{}_zha_update_device"
 SIGNAL_REMOVE_GROUP = "remove_group"
+SIGNAL_GROUP_ENTITY_REMOVED = "group_entity_removed"
 SIGNAL_GROUP_MEMBERSHIP_CHANGE = "group_membership_change"
 
 UNKNOWN = "unknown"

homeassistant/components/zha/core/device.py

@@ -561,7 +561,15 @@ async def async_add_to_group(self, group_id):
 
     async def async_remove_from_group(self, group_id):
         """Remove this device from the provided zigbee group."""
-        await self._zigpy_device.remove_from_group(group_id)
+        try:
+            await self._zigpy_device.remove_from_group(group_id)
+        except (zigpy.exceptions.DeliveryError, asyncio.TimeoutError) as ex:
+            self.debug(
+                "Failed to remove device '%s' from group: 0x%04x ex: %s",
+                self._zigpy_device.ieee,
+                group_id,
+                str(ex),
+            )
 
     async def async_bind_to_group(self, group_id, cluster_bindings):
         """Directly bind this device to a group for the given clusters."""

homeassistant/components/zha/core/discovery.py

@@ -6,7 +6,10 @@
 
 from homeassistant import const as ha_const
 from homeassistant.core import callback
-from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
 from homeassistant.helpers.entity_registry import async_entries_for_device
 from homeassistant.helpers.typing import HomeAssistantType
 
@@ -166,10 +169,30 @@ class GroupProbe:
     def __init__(self):
         """Initialize instance."""
         self._hass = None
+        self._unsubs = []
 
     def initialize(self, hass: HomeAssistantType) -> None:
         """Initialize the group probe."""
         self._hass = hass
+        self._unsubs.append(
+            async_dispatcher_connect(
+                hass, zha_const.SIGNAL_GROUP_ENTITY_REMOVED, self._reprobe_group
+            )
+        )
+
+    def cleanup(self):
+        """Clean up on when zha shuts down."""
+        for unsub in self._unsubs[:]:
+            unsub()
+            self._unsubs.remove(unsub)
+
+    def _reprobe_group(self, group_id: int) -> None:
+        """Reprobe a group for entities after its members change."""
+        zha_gateway = self._hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
+        zha_group = zha_gateway.groups.get(group_id)
+        if zha_group is None:
+            return
+        self.discover_group_entities(zha_group)
 
     @callback
     def discover_group_entities(self, group: zha_typing.ZhaGroupType) -> None:

homeassistant/components/zha/core/gateway.py

@@ -20,7 +20,10 @@
     async_get_registry as get_dev_reg,
 )
 from homeassistant.helpers.dispatcher import async_dispatcher_send
-from homeassistant.helpers.entity_registry import async_get_registry as get_ent_reg
+from homeassistant.helpers.entity_registry import (
+    async_entries_for_device,
+    async_get_registry as get_ent_reg,
+)
 
 from . import discovery, typing as zha_typing
 from .const import (
@@ -77,7 +80,7 @@
 from .device import DeviceStatus, ZHADevice
 from .group import ZHAGroup
 from .patches import apply_application_controller_patch
-from .registries import RADIO_TYPES
+from .registries import GROUP_ENTITY_DOMAINS, RADIO_TYPES
 from .store import async_get_registry
 from .typing import ZhaDeviceType, ZhaGroupType, ZigpyEndpointType, ZigpyGroupType
 
@@ -273,6 +276,9 @@ def group_member_added(
         async_dispatcher_send(
             self._hass, f"{SIGNAL_GROUP_MEMBERSHIP_CHANGE}_0x{zigpy_group.group_id:04x}"
         )
+        if len(zha_group.members) == 2:
+            # we need to do this because there wasn't already a group entity to remove and re-add
+            discovery.GROUP_PROBE.discover_group_entities(zha_group)
 
     def group_added(self, zigpy_group: ZigpyGroupType) -> None:
         """Handle zigpy group added event."""
@@ -289,6 +295,7 @@ def group_removed(self, zigpy_group: ZigpyGroupType) -> None:
         async_dispatcher_send(
             self._hass, f"{SIGNAL_REMOVE_GROUP}_0x{zigpy_group.group_id:04x}"
         )
+        self._cleanup_group_entity_registry_entries(zigpy_group)
 
     def _send_group_gateway_message(
         self, zigpy_group: ZigpyGroupType, gateway_message_type: str
@@ -368,6 +375,35 @@ def remove_entity_reference(self, entity):
                 e for e in entity_refs if e.reference_id != entity.entity_id
             ]
 
+    def _cleanup_group_entity_registry_entries(
+        self, zigpy_group: ZigpyGroupType
+    ) -> None:
+        """Remove entity registry entries for group entities when the groups are removed from HA."""
+        # first we collect the potential unique ids for entities that could be created from this group
+        possible_entity_unique_ids = [
+            f"{domain}_zha_group_0x{zigpy_group.group_id:04x}"
+            for domain in GROUP_ENTITY_DOMAINS
+        ]
+
+        # then we get all group entity entries tied to the coordinator
+        all_group_entity_entries = async_entries_for_device(
+            self.ha_entity_registry, self.coordinator_zha_device.device_id
+        )
+
+        # then we get the entity entries for this specific group by getting the entries that match
+        entries_to_remove = [
+            entry
+            for entry in all_group_entity_entries
+            if entry.unique_id in possible_entity_unique_ids
+        ]
+
+        # then we remove the entries from the entity registry
+        for entry in entries_to_remove:
+            _LOGGER.debug(
+                "cleaning up entity registry entry for entity: %s", entry.entity_id
+            )
+            self.ha_entity_registry.async_remove(entry.entity_id)
+
     @property
     def devices(self):
         """Return devices."""
@@ -557,15 +593,7 @@ async def async_create_zigpy_group(
                     )
                     tasks.append(self.devices[ieee].async_add_to_group(group_id))
                 await asyncio.gather(*tasks)
-        zha_group = self.groups.get(group_id)
-        _LOGGER.debug(
-            "Probing group: %s:0x%04x for entity discovery",
-            zha_group.name,
-            zha_group.group_id,
-        )
-        discovery.GROUP_PROBE.discover_group_entities(zha_group)
-
-        return zha_group
+        return self.groups.get(group_id)
 
     async def async_remove_zigpy_group(self, group_id: int) -> None:
         """Remove a Zigbee group from Zigpy."""

homeassistant/components/zha/entity.py

@@ -8,7 +8,10 @@
 from homeassistant.core import CALLBACK_TYPE, State, callback
 from homeassistant.helpers import entity
 from homeassistant.helpers.device_registry import CONNECTION_ZIGBEE
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
 from homeassistant.helpers.event import async_track_state_change
 from homeassistant.helpers.restore_state import RestoreEntity
 
@@ -19,6 +22,7 @@
     DATA_ZHA,
     DATA_ZHA_BRIDGE_ID,
     DOMAIN,
+    SIGNAL_GROUP_ENTITY_REMOVED,
     SIGNAL_GROUP_MEMBERSHIP_CHANGE,
     SIGNAL_REMOVE,
     SIGNAL_REMOVE_GROUP,
@@ -32,7 +36,7 @@
 RESTART_GRACE_PERIOD = 7200  # 2 hours
 
 
-class BaseZhaEntity(RestoreEntity, LogMixin, entity.Entity):
+class BaseZhaEntity(LogMixin, entity.Entity):
     """A base class for ZHA entities."""
 
     def __init__(self, unique_id: str, zha_device: ZhaDeviceType, **kwargs):
@@ -113,28 +117,11 @@ def async_update_state_attribute(self, key: str, value: Any) -> None:
     def async_set_state(self, attr_id: int, attr_name: str, value: Any) -> None:
         """Set the entity state."""
 
-    async def async_added_to_hass(self) -> None:
-        """Run when about to be added to hass."""
-        await super().async_added_to_hass()
-        self.remove_future = asyncio.Future()
-        await self.async_accept_signal(
-            None,
-            f"{SIGNAL_REMOVE}_{self.zha_device.ieee}",
-            self.async_remove,
-            signal_override=True,
-        )
-
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect entity object when removed."""
         for unsub in self._unsubs[:]:
             unsub()
             self._unsubs.remove(unsub)
-        self.zha_device.gateway.remove_entity_reference(self)
-        self.remove_future.set_result(True)
-
-    @callback
-    def async_restore_last_state(self, last_state) -> None:
-        """Restore previous state."""
 
     async def async_accept_signal(
         self, channel: ChannelType, signal: str, func: CALLABLE_T, signal_override=False
@@ -156,7 +143,7 @@ def log(self, level: int, msg: str, *args):
         _LOGGER.log(level, msg, *args)
 
 
-class ZhaEntity(BaseZhaEntity):
+class ZhaEntity(BaseZhaEntity, RestoreEntity):
     """A base class for non group ZHA entities."""
 
     def __init__(
@@ -179,6 +166,13 @@ def __init__(
     async def async_added_to_hass(self) -> None:
         """Run when about to be added to hass."""
         await super().async_added_to_hass()
+        self.remove_future = asyncio.Future()
+        await self.async_accept_signal(
+            None,
+            f"{SIGNAL_REMOVE}_{self.zha_device.ieee}",
+            self.async_remove,
+            signal_override=True,
+        )
         await self.async_check_recently_seen()
         await self.async_accept_signal(
             None,
@@ -195,6 +189,16 @@ async def async_added_to_hass(self) -> None:
             self.remove_future,
         )
 
+    async def async_will_remove_from_hass(self) -> None:
+        """Disconnect entity object when removed."""
+        await super().async_will_remove_from_hass()
+        self.zha_device.gateway.remove_entity_reference(self)
+        self.remove_future.set_result(True)
+
+    @callback
+    def async_restore_last_state(self, last_state) -> None:
+        """Restore previous state."""
+
     async def async_check_recently_seen(self) -> None:
         """Check if the device was seen within the last 2 hours."""
         last_state = await self.async_get_last_state()
@@ -244,13 +248,20 @@ async def async_added_to_hass(self) -> None:
         await self.async_accept_signal(
             None,
             f"{SIGNAL_GROUP_MEMBERSHIP_CHANGE}_0x{self._group_id:04x}",
-            self._update_group_entities,
+            self.async_remove,
             signal_override=True,
         )
 
         self._async_unsub_state_changed = async_track_state_change(
             self.hass, self._entity_ids, self.async_state_changed_listener
         )
+
+        def send_removed_signal():
+            async_dispatcher_send(
+                self.hass, SIGNAL_GROUP_ENTITY_REMOVED, self._group_id
+            )
+
+        self.async_on_remove(send_removed_signal)
         await self.async_update()
 
     @callback
@@ -260,17 +271,6 @@ def async_state_changed_listener(
         """Handle child updates."""
         self.async_schedule_update_ha_state(True)
 
-    def _update_group_entities(self):
-        """Update tracked entities when membership changes."""
-        group = self.zha_device.gateway.get_group(self._group_id)
-        self._entity_ids = group.get_domain_entity_ids(self.platform.domain)
-        if self._async_unsub_state_changed is not None:
-            self._async_unsub_state_changed()
-
-        self._async_unsub_state_changed = async_track_state_change(
-            self.hass, self._entity_ids, self.async_state_changed_listener
-        )
-
     async def async_will_remove_from_hass(self) -> None:
         """Handle removal from Home Assistant."""
         await super().async_will_remove_from_hass()

tests/components/zha/test_light.py

@@ -539,3 +539,21 @@ async def async_test_zha_group_light_entity(
     await zha_group.async_add_members([device_light_3.ieee])
     await dev3_cluster_on_off.on()
     assert hass.states.get(entity_id).state == STATE_ON
+
+    # make the group have only 1 member and now there should be no entity
+    await zha_group.async_remove_members([device_light_2.ieee, device_light_3.ieee])
+    assert len(zha_group.members) == 1
+    assert hass.states.get(entity_id).state is None
+    # make sure the entity registry entry is still there
+    assert zha_gateway.ha_entity_registry.async_get(entity_id) is not None
+
+    # add a member back and ensure that the group entity was created again
+    await zha_group.async_add_members([device_light_3.ieee])
+    await dev3_cluster_on_off.on()
+    assert hass.states.get(entity_id).state == STATE_ON
+
+    # remove the group and ensure that there is no entity and that the entity registry is cleaned up
+    assert zha_gateway.ha_entity_registry.async_get(entity_id) is not None
+    await zha_gateway.async_remove_zigpy_group(zha_group.group_id)
+    assert hass.states.get(entity_id).state is None
+    assert zha_gateway.ha_entity_registry.async_get(entity_id) is None