From 1282370ccbe10b3fe18f5d47090cf3a830985ddf Mon Sep 17 00:00:00 2001
From: Charles Garwood <cgarwood@gmail.com>
Date: Sun, 19 May 2019 05:14:11 -0400
Subject: [PATCH] Entity Cleanup on Z-Wave node removal (#23633)

* Initial groundwork for entity cleanup on node removal

* Connect node_removed to dispatcher

* update docstring

* Add node_removal test

* Address review comments

* Use hass.add_job instead of run_coroutine_threadsafe
---
 homeassistant/components/zwave/__init__.py    | 21 ++++++++++
 homeassistant/components/zwave/node_entity.py | 11 +++++
 tests/components/zwave/test_init.py           | 42 +++++++++++++++++++
 3 files changed, 74 insertions(+)

diff --git a/homeassistant/components/zwave/__init__.py b/homeassistant/components/zwave/__init__.py
index 741c6f852a8..10046825ad3 100644
--- a/homeassistant/components/zwave/__init__.py
+++ b/homeassistant/components/zwave/__init__.py
@@ -376,6 +376,25 @@ async def async_setup_entry(hass, config_entry):
         hass.add_job(check_has_unique_id, entity, _on_ready, _on_timeout,
                      hass.loop)
 
+    def node_removed(node):
+        node_id = node.node_id
+        node_key = 'node-{}'.format(node_id)
+        _LOGGER.info("Node Removed: %s",
+                     hass.data[DATA_DEVICES][node_key])
+        for key in list(hass.data[DATA_DEVICES]):
+            if not key.startswith('{}-'.format(node_id)):
+                continue
+
+            entity = hass.data[DATA_DEVICES][key]
+            _LOGGER.info('Removing Entity - value: %s - entity_id: %s',
+                         key, entity.entity_id)
+            hass.add_job(entity.node_removed())
+            del hass.data[DATA_DEVICES][key]
+
+        entity = hass.data[DATA_DEVICES][node_key]
+        hass.add_job(entity.node_removed())
+        del hass.data[DATA_DEVICES][node_key]
+
     def network_ready():
         """Handle the query of all awake nodes."""
         _LOGGER.info("Z-Wave network is ready for use. All awake nodes "
@@ -399,6 +418,8 @@ async def async_setup_entry(hass, config_entry):
         value_added, ZWaveNetwork.SIGNAL_VALUE_ADDED, weak=False)
     dispatcher.connect(
         node_added, ZWaveNetwork.SIGNAL_NODE_ADDED, weak=False)
+    dispatcher.connect(
+        node_removed, ZWaveNetwork.SIGNAL_NODE_REMOVED, weak=False)
     dispatcher.connect(
         network_ready, ZWaveNetwork.SIGNAL_AWAKE_NODES_QUERIED, weak=False)
     dispatcher.connect(
diff --git a/homeassistant/components/zwave/node_entity.py b/homeassistant/components/zwave/node_entity.py
index 2339b8aba36..0a24f888c20 100644
--- a/homeassistant/components/zwave/node_entity.py
+++ b/homeassistant/components/zwave/node_entity.py
@@ -3,6 +3,7 @@ import logging
 
 from homeassistant.core import callback
 from homeassistant.const import ATTR_BATTERY_LEVEL, ATTR_WAKEUP, ATTR_ENTITY_ID
+from homeassistant.helpers.entity_registry import async_get_registry
 from homeassistant.helpers.entity import Entity
 
 from .const import (
@@ -74,6 +75,16 @@ class ZWaveBaseEntity(Entity):
         if self.hass and self.platform:
             self.hass.add_job(_async_remove_and_add)
 
+    async def node_removed(self):
+        """Call when a node is removed from the Z-Wave network."""
+        await self.async_remove()
+
+        registry = await async_get_registry(self.hass)
+        if self.entity_id not in registry.entities:
+            return
+
+        registry.async_remove(self.entity_id)
+
 
 class ZWaveNodeEntity(ZWaveBaseEntity):
     """Representation of a Z-Wave node."""
diff --git a/tests/components/zwave/test_init.py b/tests/components/zwave/test_init.py
index 7fc9f55cf03..69ee7c45a9b 100644
--- a/tests/components/zwave/test_init.py
+++ b/tests/components/zwave/test_init.py
@@ -226,6 +226,48 @@ async def test_device_entity(hass, mock_openzwave):
     assert device.device_state_attributes[zwave.ATTR_POWER] == 50.123
 
 
+async def test_node_removed(hass, mock_openzwave):
+    """Test node removed in base class."""
+    # Create a mock node & node entity
+    node = MockNode(node_id='10', name='Mock Node')
+    value = MockValue(data=False, node=node, instance=2, object_id='11',
+                      label='Sensor',
+                      command_class=const.COMMAND_CLASS_SENSOR_BINARY)
+    power_value = MockValue(data=50.123456, node=node, precision=3,
+                            command_class=const.COMMAND_CLASS_METER)
+    values = MockEntityValues(primary=value, power=power_value)
+    device = zwave.ZWaveDeviceEntity(values, 'zwave')
+    device.hass = hass
+    device.entity_id = 'zwave.mock_node'
+    device.value_added()
+    device.update_properties()
+    await hass.async_block_till_done()
+
+    # Save it to the entity registry
+    registry = mock_registry(hass)
+    registry.async_get_or_create('zwave', 'zwave', device.unique_id)
+    device.entity_id = registry.async_get_entity_id(
+        'zwave', 'zwave', device.unique_id)
+
+    # Create dummy entity registry entries for other integrations
+    hue_entity = registry.async_get_or_create('light', 'hue', 1234)
+    zha_entity = registry.async_get_or_create('sensor', 'zha', 5678)
+
+    # Verify our Z-Wave entity is registered
+    assert registry.async_is_registered(device.entity_id)
+
+    # Remove it
+    entity_id = device.entity_id
+    await device.node_removed()
+
+    # Verify registry entry for our Z-Wave node is gone
+    assert not registry.async_is_registered(entity_id)
+
+    # Verify registry entries for our other entities remain
+    assert registry.async_is_registered(hue_entity.entity_id)
+    assert registry.async_is_registered(zha_entity.entity_id)
+
+
 async def test_node_discovery(hass, mock_openzwave):
     """Test discovery of a node."""
     mock_receivers = []