Fix removing nodes in zwave_js integration (#45676)

homeassistant/components/zwave_js/__init__.py

@@ -1,6 +1,7 @@
 """The Z-Wave JS integration."""
 import asyncio
 import logging
+from typing import Tuple
 
 from async_timeout import timeout
 from zwave_js_server.client import Client as ZwaveClient
@@ -36,6 +37,12 @@ async def async_setup(hass: HomeAssistant, config: dict) -> bool:
     return True
 
 
+@callback
+def get_device_id(client: ZwaveClient, node: ZwaveNode) -> Tuple[str, str]:
+    """Get device registry identifier for Z-Wave node."""
+    return (DOMAIN, f"{client.driver.controller.home_id}-{node.node_id}")
+
+
 @callback
 def register_node_in_dev_reg(
     hass: HomeAssistant,
@@ -47,7 +54,7 @@ def register_node_in_dev_reg(
     """Register node in dev reg."""
     device = dev_reg.async_get_or_create(
         config_entry_id=entry.entry_id,
-        identifiers={(DOMAIN, f"{client.driver.controller.home_id}-{node.node_id}")},
+        identifiers={get_device_id(client, node)},
         sw_version=node.firmware_version,
         name=node.name or node.device_config.description or f"Node {node.node_id}",
         model=node.device_config.label,
@@ -118,6 +125,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         # some visual feedback that something is (in the process of) being added
         register_node_in_dev_reg(hass, entry, dev_reg, client, node)
 
+    @callback
+    def async_on_node_removed(node: ZwaveNode) -> None:
+        """Handle node removed event."""
+        # grab device in device registry attached to this node
+        dev_id = get_device_id(client, node)
+        device = dev_reg.async_get_device({dev_id})
+        # note: removal of entity registry is handled by core
+        dev_reg.async_remove_device(device.id)
+
     async def handle_ha_shutdown(event: Event) -> None:
         """Handle HA shutdown."""
         await client.disconnect()
@@ -171,6 +187,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         client.driver.controller.on(
             "node added", lambda event: async_on_node_added(event["node"])
         )
+        # listen for nodes being removed from the mesh
+        # NOTE: This will not remove nodes that were removed when HA was not running
+        client.driver.controller.on(
+            "node removed", lambda event: async_on_node_removed(event["node"])
+        )
 
     hass.async_create_task(start_platforms())
 

homeassistant/components/zwave_js/api.py

@@ -5,15 +5,12 @@ import logging
 from aiohttp import hdrs, web, web_exceptions
 import voluptuous as vol
 from zwave_js_server import dump
-from zwave_js_server.client import Client as ZwaveClient
-from zwave_js_server.model.node import Node as ZwaveNode
 
 from homeassistant.components import websocket_api
 from homeassistant.components.http.view import HomeAssistantView
 from homeassistant.components.websocket_api.connection import ActiveConnection
 from homeassistant.const import CONF_URL
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import device_registry
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.device_registry import DeviceEntry
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -248,9 +245,6 @@ async def websocket_remove_node(
             "node_id": node.node_id,
         }
 
-        # Remove from device registry
-        hass.async_create_task(remove_from_device_registry(hass, client, node))
-
         connection.send_message(
             websocket_api.event_message(
                 msg[ID], {"event": "node removed", "node": node_details}
@@ -272,20 +266,6 @@ async def websocket_remove_node(
     )
 
 
-async def remove_from_device_registry(
-    hass: HomeAssistant, client: ZwaveClient, node: ZwaveNode
-) -> None:
-    """Remove a node from the device registry."""
-    registry = await device_registry.async_get_registry(hass)
-    device = registry.async_get_device(
-        {(DOMAIN, f"{client.driver.controller.home_id}-{node.node_id}")}
-    )
-    if device is None:
-        return
-
-    registry.async_remove_device(device.id)
-
-
 class DumpView(HomeAssistantView):
     """View to dump the state of the Z-Wave JS server."""