diff --git a/homeassistant/components/zha/api.py b/homeassistant/components/zha/api.py index 077863d34b8..0604c2fada4 100644 --- a/homeassistant/components/zha/api.py +++ b/homeassistant/components/zha/api.py @@ -139,7 +139,7 @@ def async_get_device_info(hass, device, ha_device_registry=None): ret_device = {} ret_device.update(device.device_info) ret_device['entities'] = [{ - 'entity_id': entity_ref.entity.entity_id, + 'entity_id': entity_ref.reference_id, NAME: entity_ref.device_info[NAME] } for entity_ref in zha_gateway.device_registry[device.ieee]] diff --git a/homeassistant/components/zha/core/const.py b/homeassistant/components/zha/core/const.py index 55001a0a419..23b2bb99050 100644 --- a/homeassistant/components/zha/core/const.py +++ b/homeassistant/components/zha/core/const.py @@ -103,6 +103,7 @@ SIGNAL_MOVE_LEVEL = "move_level" SIGNAL_SET_LEVEL = "set_level" SIGNAL_STATE_ATTR = "update_state_attribute" SIGNAL_AVAILABLE = 'available' +SIGNAL_REMOVE = 'remove' QUIRK_APPLIED = 'quirk_applied' QUIRK_CLASS = 'quirk_class' diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py index 8a94f0928b6..307d85a8d9e 100644 --- a/homeassistant/components/zha/core/gateway.py +++ b/homeassistant/components/zha/core/gateway.py @@ -27,8 +27,8 @@ from .const import ( DEBUG_LEVELS, DEFAULT_BAUDRATE, DEFAULT_DATABASE_NAME, DEVICE_FULL_INIT, DEVICE_INFO, DEVICE_JOINED, DEVICE_REMOVED, DOMAIN, IEEE, LOG_ENTRY, LOG_OUTPUT, MODEL, NWK, ORIGINAL, RADIO, RADIO_DESCRIPTION, RAW_INIT, - SIGNATURE, TYPE, UNKNOWN_MANUFACTURER, UNKNOWN_MODEL, ZHA, ZHA_GW_MSG, - ZIGPY, ZIGPY_DECONZ, ZIGPY_XBEE) + SIGNAL_REMOVE, SIGNATURE, TYPE, UNKNOWN_MANUFACTURER, UNKNOWN_MODEL, ZHA, + ZHA_GW_MSG, ZIGPY, ZIGPY_DECONZ, ZIGPY_XBEE) from .device import DeviceStatus, ZHADevice from .discovery import ( async_create_device_entity, async_dispatch_discovery_info, @@ -41,8 +41,7 @@ _LOGGER = logging.getLogger(__name__) EntityReference = collections.namedtuple( 'EntityReference', - 'entity device_info' -) + 'reference_id zha_device cluster_channels device_info remove_future') class ZHAGateway: @@ -149,8 +148,8 @@ class ZHAGateway: if entity_refs is not None: remove_tasks = [] for entity_ref in entity_refs: - remove_tasks.append(entity_ref.entity.async_remove()) - await asyncio.gather(*remove_tasks) + remove_tasks.append(entity_ref.remove_future) + await asyncio.wait(remove_tasks) ha_device_registry = await get_dev_reg(self._hass) reg_device = ha_device_registry.async_get_device( {(DOMAIN, str(device.ieee))}, set()) @@ -164,6 +163,10 @@ class ZHAGateway: if zha_device is not None: device_info = async_get_device_info(self._hass, zha_device) zha_device.async_unsub_dispatcher() + async_dispatcher_send( + self._hass, + "{}_{}".format(SIGNAL_REMOVE, str(zha_device.ieee)) + ) asyncio.ensure_future( self._async_remove_device(zha_device, entity_refs) ) @@ -185,7 +188,7 @@ class ZHAGateway: """Return entity reference for given entity_id if found.""" for entity_reference in itertools.chain.from_iterable( self.device_registry.values()): - if entity_id == entity_reference.entity.entity_id: + if entity_id == entity_reference.reference_id: return entity_reference @property @@ -198,10 +201,18 @@ class ZHAGateway: """Return entities by ieee.""" return self._device_registry - def register_entity_reference(self, ieee, entity, device_info): + def register_entity_reference( + self, ieee, reference_id, zha_device, cluster_channels, + device_info, remove_future): """Record the creation of a hass entity associated with ieee.""" self._device_registry[ieee].append( - EntityReference(entity=entity, device_info=device_info) + EntityReference( + reference_id=reference_id, + zha_device=zha_device, + cluster_channels=cluster_channels, + device_info=device_info, + remove_future=remove_future + ) ) @callback diff --git a/homeassistant/components/zha/entity.py b/homeassistant/components/zha/entity.py index 9e4c583875f..a854a5c9a6e 100644 --- a/homeassistant/components/zha/entity.py +++ b/homeassistant/components/zha/entity.py @@ -1,5 +1,6 @@ """Entity for Zigbee Home Automation.""" +import asyncio import logging import time @@ -11,7 +12,8 @@ from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.util import slugify from .core.const import ( - ATTR_MANUFACTURER, DATA_ZHA, DATA_ZHA_BRIDGE_ID, DOMAIN, MODEL, NAME) + ATTR_MANUFACTURER, DATA_ZHA, DATA_ZHA_BRIDGE_ID, DOMAIN, MODEL, NAME, + SIGNAL_REMOVE) _LOGGER = logging.getLogger(__name__) @@ -48,6 +50,7 @@ class ZhaEntity(RestoreEntity, entity.Entity): self._available = False self._component = kwargs['component'] self._unsubs = [] + self.remove_future = asyncio.Future() for channel in channels: self.cluster_channels[channel.name] = channel @@ -125,9 +128,14 @@ class ZhaEntity(RestoreEntity, entity.Entity): None, "{}_{}".format(self.zha_device.available_signal, 'entity'), self.async_set_available, signal_override=True) - self._zha_device.gateway.register_entity_reference( - self._zha_device.ieee, self, self.device_info + await self.async_accept_signal( + None, "{}_{}".format(SIGNAL_REMOVE, str(self.zha_device.ieee)), + self.async_remove, + signal_override=True ) + self._zha_device.gateway.register_entity_reference( + self._zha_device.ieee, self.entity_id, self._zha_device, + self.cluster_channels, self.device_info, self.remove_future) async def async_check_recently_seen(self): """Check if the device was seen within the last 2 hours.""" @@ -145,6 +153,7 @@ class ZhaEntity(RestoreEntity, entity.Entity): """Disconnect entity object when removed.""" for unsub in self._unsubs: unsub() + self.remove_future.set_result(True) @callback def async_restore_last_state(self, last_state):