diff --git a/homeassistant/components/zha/core/device.py b/homeassistant/components/zha/core/device.py index dcb4fe7ca0e..835b9ee7e81 100644 --- a/homeassistant/components/zha/core/device.py +++ b/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() diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py index d1ccaf8265c..c77c9fbbb7b 100644 --- a/homeassistant/components/zha/core/gateway.py +++ b/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)) diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index 13a013522fb..77c788035ab 100644 --- a/homeassistant/helpers/device_registry.py +++ b/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: diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 2fb32d5214e..302eda79426 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/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):