From 13dda7bd9859835e25412c3596b25dfc9f6fb334 Mon Sep 17 00:00:00 2001 From: "David F. Mulcahey" Date: Fri, 10 Apr 2020 16:17:48 -0400 Subject: [PATCH] Cleanup ZHA group entity lifecycle (#33977) * Clean up ZHA group entity lifecycle * group entities don't use state restore * add tests --- homeassistant/components/zha/__init__.py | 2 + homeassistant/components/zha/core/const.py | 1 + homeassistant/components/zha/core/device.py | 10 ++- .../components/zha/core/discovery.py | 25 +++++++- homeassistant/components/zha/core/gateway.py | 50 +++++++++++---- homeassistant/components/zha/entity.py | 64 +++++++++---------- tests/components/zha/test_light.py | 18 ++++++ 7 files changed, 125 insertions(+), 45 deletions(-) diff --git a/homeassistant/components/zha/__init__.py b/homeassistant/components/zha/__init__.py index 2af35e8fb92..9e59b63adb4 100644 --- a/homeassistant/components/zha/__init__.py +++ b/homeassistant/components/zha/__init__.py @@ -32,6 +32,7 @@ from .core.const import ( 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, []) diff --git a/homeassistant/components/zha/core/const.py b/homeassistant/components/zha/core/const.py index 055b5627bb1..b181b848f04 100644 --- a/homeassistant/components/zha/core/const.py +++ b/homeassistant/components/zha/core/const.py @@ -213,6 +213,7 @@ SIGNAL_SET_LEVEL = "set_level" 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" diff --git a/homeassistant/components/zha/core/device.py b/homeassistant/components/zha/core/device.py index 287ad0dd522..5782ce23083 100644 --- a/homeassistant/components/zha/core/device.py +++ b/homeassistant/components/zha/core/device.py @@ -561,7 +561,15 @@ class ZHADevice(LogMixin): 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.""" diff --git a/homeassistant/components/zha/core/discovery.py b/homeassistant/components/zha/core/discovery.py index 90ec0e6e250..4540c9158de 100644 --- a/homeassistant/components/zha/core/discovery.py +++ b/homeassistant/components/zha/core/discovery.py @@ -6,7 +6,10 @@ from typing import Callable, List, Tuple 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: diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py index e032de4d94c..e97e2185dc5 100644 --- a/homeassistant/components/zha/core/gateway.py +++ b/homeassistant/components/zha/core/gateway.py @@ -20,7 +20,10 @@ from homeassistant.helpers.device_registry import ( 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 .const import ( 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 @@ class ZHAGateway: 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 @@ class ZHAGateway: 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 @@ class ZHAGateway: 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 @@ class ZHAGateway: ) 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.""" diff --git a/homeassistant/components/zha/entity.py b/homeassistant/components/zha/entity.py index dd12924f4b0..fe213f7920b 100644 --- a/homeassistant/components/zha/entity.py +++ b/homeassistant/components/zha/entity.py @@ -8,7 +8,10 @@ from typing import Any, Awaitable, Dict, List, Optional 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 @@ from .core.const import ( DATA_ZHA, DATA_ZHA_BRIDGE_ID, DOMAIN, + SIGNAL_GROUP_ENTITY_REMOVED, SIGNAL_GROUP_MEMBERSHIP_CHANGE, SIGNAL_REMOVE, SIGNAL_REMOVE_GROUP, @@ -32,7 +36,7 @@ ENTITY_SUFFIX = "entity_suffix" 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 @@ class BaseZhaEntity(RestoreEntity, LogMixin, entity.Entity): 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 @@ class BaseZhaEntity(RestoreEntity, LogMixin, entity.Entity): _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 @@ class ZhaEntity(BaseZhaEntity): 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 @@ class ZhaEntity(BaseZhaEntity): 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 @@ class ZhaGroupEntity(BaseZhaEntity): 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 @@ class ZhaGroupEntity(BaseZhaEntity): """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() diff --git a/tests/components/zha/test_light.py b/tests/components/zha/test_light.py index 9bdd4966a4a..f297bfa5bf0 100644 --- a/tests/components/zha/test_light.py +++ b/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