go back to signals and no hard entity references (#24894)

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

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'

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

homeassistant/components/zha/entity.py

 """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):