Cleanup ZHA group entity lifecycle (#33977)

* 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

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