From 7a8130cd2ba2d96c467976b7b6e79b8b3dfd7605 Mon Sep 17 00:00:00 2001
From: "David F. Mulcahey" <david.mulcahey@me.com>
Date: Sun, 21 Jul 2019 13:00:27 -0400
Subject: [PATCH] fix remove and re-add scenario (#25370)

---
 homeassistant/components/zha/core/gateway.py | 7 +++++++
 homeassistant/components/zha/entity.py       | 7 +++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py
index 351ad1c5a67..20bf30dec87 100644
--- a/homeassistant/components/zha/core/gateway.py
+++ b/homeassistant/components/zha/core/gateway.py
@@ -190,6 +190,13 @@ class ZHAGateway:
             if entity_id == entity_reference.reference_id:
                 return entity_reference
 
+    def remove_entity_reference(self, entity):
+        """Remove entity reference for given entity_id if found."""
+        if entity.zha_device.ieee in self.device_registry:
+            entity_refs = self.device_registry.get(entity.zha_device.ieee)
+            self.device_registry[entity.zha_device.ieee] = [
+                e for e in entity_refs if e.reference_id != entity.entity_id]
+
     @property
     def devices(self):
         """Return devices."""
diff --git a/homeassistant/components/zha/entity.py b/homeassistant/components/zha/entity.py
index a854a5c9a6e..77cb5b8deda 100644
--- a/homeassistant/components/zha/entity.py
+++ b/homeassistant/components/zha/entity.py
@@ -50,7 +50,7 @@ class ZhaEntity(RestoreEntity, entity.Entity):
         self._available = False
         self._component = kwargs['component']
         self._unsubs = []
-        self.remove_future = asyncio.Future()
+        self.remove_future = None
         for channel in channels:
             self.cluster_channels[channel.name] = channel
 
@@ -123,6 +123,7 @@ class ZhaEntity(RestoreEntity, entity.Entity):
     async def async_added_to_hass(self):
         """Run when about to be added to hass."""
         await super().async_added_to_hass()
+        self.remove_future = asyncio.Future()
         await self.async_check_recently_seen()
         await self.async_accept_signal(
             None, "{}_{}".format(self.zha_device.available_signal, 'entity'),
@@ -151,8 +152,10 @@ class ZhaEntity(RestoreEntity, entity.Entity):
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect entity object when removed."""
-        for unsub in self._unsubs:
+        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