evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Cleanup ZHA group entity lifecycle (#33977)

Browse source 
* Clean up ZHA group entity lifecycle

* group entities don't use state restore

* add tests
This commit is contained in:
David F. Mulcahey 2020-04-10 16:17:48 -04:00 • committed by GitHub
parent 14c35c9223
commit 13dda7bd98
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 125 additions and 45 deletions
Download patch file Download diff file Expand all files Collapse all files

2
homeassistant/components/zha/__init__.py
View file

@ -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, [])

1
homeassistant/components/zha/core/const.py
View file

@ -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"

10
homeassistant/components/zha/core/device.py
View file

@ -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."""

25
homeassistant/components/zha/core/discovery.py
View file

@ -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:

50
homeassistant/components/zha/core/gateway.py
View file

@ -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."""

64
homeassistant/components/zha/entity.py
View file

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

18
tests/components/zha/test_light.py
View file

@ -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