Remove device and entity registry entries when removing a ZHA device (#24369)

* cleanup when device is removed

fixes

* cleanup

homeassistant/components/zha/core/device.py

@@ -312,7 +312,8 @@ class ZHADevice:
                 ex
             )
 
-    async def async_unsub_dispatcher(self):
+    @callback
+    def async_unsub_dispatcher(self):
         """Unsubscribe the dispatcher."""
         if self._unsub:
             self._unsub()

homeassistant/components/zha/core/gateway.py

@@ -14,6 +14,8 @@ import traceback
 
 from homeassistant.components.system_log import LogEntry, _figure_out_source
 from homeassistant.core import callback
+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_component import EntityComponent
 
@@ -146,13 +148,20 @@ class ZHAGateway:
         """Handle device leaving the network."""
         pass
 
+    async def _async_remove_device(self, device):
+        ha_device_registry = await get_dev_reg(self._hass)
+        reg_device = ha_device_registry.async_get_device(
+            {(DOMAIN, str(device.ieee))}, set())
+        ha_device_registry.async_remove_device(reg_device.id)
+
     def device_removed(self, device):
         """Handle device being removed from the network."""
         zha_device = self._devices.pop(device.ieee, None)
         self._device_registry.pop(device.ieee, None)
         if zha_device is not None:
             device_info = async_get_device_info(self._hass, zha_device)
-            self._hass.async_create_task(zha_device.async_unsub_dispatcher())
+            zha_device.async_unsub_dispatcher()
+            asyncio.ensure_future(self._async_remove_device(zha_device))
             async_dispatcher_send(
                 self._hass,
                 "{}_{}".format(SIGNAL_REMOVE, str(zha_device.ieee))

homeassistant/helpers/device_registry.py

@@ -219,7 +219,8 @@ class DeviceRegistry:
 
         return new
 
-    def _async_remove_device(self, device_id):
+    def async_remove_device(self, device_id):
+        """Remove a device from the device registry."""
         del self.devices[device_id]
         self.hass.bus.async_fire(EVENT_DEVICE_REGISTRY_UPDATED, {
             'action': 'remove',
@@ -297,7 +298,7 @@ class DeviceRegistry:
                 self._async_update_device(
                     dev_id, remove_config_entry_id=config_entry_id)
         for dev_id in remove:
-            self._async_remove_device(dev_id)
+            self.async_remove_device(dev_id)
 
     @callback
     def async_clear_area_id(self, area_id: str) -> None:

homeassistant/helpers/entity_registry.py

@@ -17,6 +17,7 @@ import weakref
 import attr
 
 from homeassistant.core import callback, split_entity_id, valid_entity_id
+from homeassistant.helpers.device_registry import EVENT_DEVICE_REGISTRY_UPDATED
 from homeassistant.loader import bind_hass
 from homeassistant.util import ensure_unique_string, slugify
 from homeassistant.util.yaml import load_yaml
@@ -84,6 +85,10 @@ class EntityRegistry:
         self.hass = hass
         self.entities = None
         self._store = hass.helpers.storage.Store(STORAGE_VERSION, STORAGE_KEY)
+        self.hass.bus.async_listen(
+            EVENT_DEVICE_REGISTRY_UPDATED,
+            self.async_device_removed
+        )
 
     @callback
     def async_is_registered(self, entity_id):
@@ -169,6 +174,19 @@ class EntityRegistry:
         })
         self.async_schedule_save()
 
+    @callback
+    def async_device_removed(self, event):
+        """Handle the removal of a device.
+
+        Remove entities from the registry that are associated to a device when
+        the device is removed.
+        """
+        if event.data['action'] != 'remove':
+            return
+        entities = async_entries_for_device(self, event.data['device_id'])
+        for entity in entities:
+            self.async_remove(entity.entity_id)
+
     @callback
     def async_update_entity(self, entity_id, *, name=_UNDEF,
                             new_entity_id=_UNDEF, new_unique_id=_UNDEF):