Add add & remove node commands to zwave_js websocket api

homeassistant/components/zwave_js/__init__.py

@@ -14,7 +14,13 @@
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
-from .const import DATA_CLIENT, DATA_UNSUBSCRIBE, DOMAIN, PLATFORMS
+from .const import (
+    DATA_CLIENT,
+    DATA_UNSUBSCRIBE,
+    DOMAIN,
+    EVENT_DEVICE_ADDED_TO_REGISTRY,
+    PLATFORMS,
+)
 from .discovery import async_discover_values
 from .websocket_api import async_register_api
 
@@ -30,13 +36,14 @@ async def async_setup(hass: HomeAssistant, config: dict) -> bool:
 
 @callback
 def register_node_in_dev_reg(
+    hass: HomeAssistant,
     entry: ConfigEntry,
     dev_reg: device_registry.DeviceRegistry,
     client: ZwaveClient,
     node: ZwaveNode,
 ) -> None:
     """Register node in dev reg."""
-    dev_reg.async_get_or_create(
+    device = dev_reg.async_get_or_create(
         config_entry_id=entry.entry_id,
         identifiers={(DOMAIN, f"{client.driver.controller.home_id}-{node.node_id}")},
         sw_version=node.firmware_version,
@@ -45,6 +52,8 @@ def register_node_in_dev_reg(
         manufacturer=node.device_config.manufacturer,
     )
 
+    async_dispatcher_send(hass, EVENT_DEVICE_ADDED_TO_REGISTRY, device)
+
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Z-Wave JS from a config entry."""
@@ -75,7 +84,7 @@ def async_on_node_ready(node: ZwaveNode) -> None:
         LOGGER.debug("Processing node %s", node)
 
         # register (or update) node in device registry
-        register_node_in_dev_reg(entry, dev_reg, client, node)
+        register_node_in_dev_reg(hass, entry, dev_reg, client, node)
 
         # run discovery on all node values and create/update entities
         for disc_info in async_discover_values(node):
@@ -98,7 +107,7 @@ def async_on_node_added(node: ZwaveNode) -> None:
         )
         # we do submit the node to device registry so user has
         # some visual feedback that something is (in the process of) being added
-        register_node_in_dev_reg(entry, dev_reg, client, node)
+        register_node_in_dev_reg(hass, entry, dev_reg, client, node)
 
     async def handle_ha_shutdown(event: Event) -> None:
         """Handle HA shutdown."""

homeassistant/components/zwave_js/const.py

@@ -7,3 +7,5 @@
 
 DATA_CLIENT = "client"
 DATA_UNSUBSCRIBE = "unsubs"
+
+EVENT_DEVICE_ADDED_TO_REGISTRY = f"{DOMAIN}_device_added_to_registry"

homeassistant/components/zwave_js/websocket_api.py

@@ -3,12 +3,17 @@
 import logging
 
 import voluptuous as vol
+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.websocket_api.connection import ActiveConnection
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import device_registry
+from homeassistant.helpers.device_registry import DeviceEntry
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
-from .const import DATA_CLIENT, DOMAIN
+from .const import DATA_CLIENT, DOMAIN, EVENT_DEVICE_ADDED_TO_REGISTRY
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -21,6 +26,10 @@
 def async_register_api(hass: HomeAssistant) -> None:
     """Register all of our api endpoints."""
     websocket_api.async_register_command(hass, websocket_network_status)
+    websocket_api.async_register_command(hass, websocket_add_node)
+    websocket_api.async_register_command(hass, websocket_stop_inclusion)
+    websocket_api.async_register_command(hass, websocket_remove_node)
+    websocket_api.async_register_command(hass, websocket_stop_exclusion)
 
 
 @websocket_api.require_admin
@@ -50,3 +59,191 @@ def websocket_network_status(
         msg[ID],
         data,
     )
+
+
+@websocket_api.require_admin  # type: ignore
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zwave_js/add_node",
+        vol.Required(ENTRY_ID): str,
+        vol.Optional("secure", default=False): bool,
+    }
+)
+async def websocket_add_node(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+) -> None:
+    """Add a node to the Z-Wave network."""
+    entry_id = msg[ENTRY_ID]
+    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
+    controller = client.driver.controller
+    include_non_secure = not msg["secure"]
+
+    @callback
+    def async_cleanup() -> None:
+        """Remove signal listeners."""
+        for unsub in unsubs:
+            unsub()
+
+    @callback
+    def forward_event(event: dict) -> None:
+        connection.send_message(
+            websocket_api.event_message(msg[ID], {"event": event["event"]})
+        )
+
+    @callback
+    def node_added(event: dict) -> None:
+        node = event["node"]
+        node_details = {
+            "node_id": node.node_id,
+            "status": node.status,
+            "ready": node.ready,
+        }
+        connection.send_message(
+            websocket_api.event_message(
+                msg[ID], {"event": "node added", "node": node_details}
+            )
+        )
+
+    @callback
+    def device_registered(device: DeviceEntry) -> None:
+        device_details = {"name": device.name, "id": device.id}
+        connection.send_message(
+            websocket_api.event_message(
+                msg[ID], {"event": "device registered", "device": device_details}
+            )
+        )
+
+    connection.subscriptions[msg["id"]] = async_cleanup
+    unsubs = [
+        controller.on("inclusion started", forward_event),
+        controller.on("inclusion failed", forward_event),
+        controller.on("inclusion stopped", forward_event),
+        controller.on("node added", node_added),
+        async_dispatcher_connect(
+            hass, EVENT_DEVICE_ADDED_TO_REGISTRY, device_registered
+        ),
+    ]
+
+    result = await controller.async_begin_inclusion(include_non_secure)
+    connection.send_result(
+        msg[ID],
+        result,
+    )
+
+
+@websocket_api.require_admin  # type: ignore
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zwave_js/stop_inclusion",
+        vol.Required(ENTRY_ID): str,
+    }
+)
+async def websocket_stop_inclusion(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+) -> None:
+    """Cancel adding a node to the Z-Wave network."""
+    entry_id = msg[ENTRY_ID]
+    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
+    controller = client.driver.controller
+    result = await controller.async_stop_inclusion()
+    connection.send_result(
+        msg[ID],
+        result,
+    )
+
+
+@websocket_api.require_admin  # type: ignore
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zwave_js/stop_exclusion",
+        vol.Required(ENTRY_ID): str,
+    }
+)
+async def websocket_stop_exclusion(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+) -> None:
+    """Cancel removing a node from the Z-Wave network."""
+    entry_id = msg[ENTRY_ID]
+    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
+    controller = client.driver.controller
+    result = await controller.async_stop_exclusion()
+    connection.send_result(
+        msg[ID],
+        result,
+    )
+
+
+@websocket_api.require_admin  # type:ignore
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zwave_js/remove_node",
+        vol.Required(ENTRY_ID): str,
+    }
+)
+async def websocket_remove_node(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+) -> None:
+    """Remove a node from the Z-Wave network."""
+    entry_id = msg[ENTRY_ID]
+    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
+    controller = client.driver.controller
+
+    @callback
+    def async_cleanup() -> None:
+        """Remove signal listeners."""
+        for unsub in unsubs:
+            unsub()
+
+    @callback
+    def forward_event(event: dict) -> None:
+        connection.send_message(
+            websocket_api.event_message(msg[ID], {"event": event["event"]})
+        )
+
+    @callback
+    def node_removed(event: dict) -> None:
+        node = event["node"]
+        node_details = {
+            "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}
+            )
+        )
+
+    connection.subscriptions[msg["id"]] = async_cleanup
+    unsubs = [
+        controller.on("exclusion started", forward_event),
+        controller.on("exclusion failed", forward_event),
+        controller.on("exclusion stopped", forward_event),
+        controller.on("node removed", node_removed),
+    ]
+
+    result = await controller.async_begin_exclusion()
+    connection.send_result(
+        msg[ID],
+        result,
+    )
+
+
+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)

tests/components/zwave_js/conftest.py

@@ -3,6 +3,7 @@
 from unittest.mock import DEFAULT, patch
 
 import pytest
+from zwave_js_server.event import Event
 from zwave_js_server.model.driver import Driver
 from zwave_js_server.model.node import Node
 from zwave_js_server.version import VersionInfo
@@ -75,6 +76,12 @@ def climate_radio_thermostat_ct100_plus_state_fixture():
     )
 
 
+@pytest.fixture(name="nortek_thermostat_state", scope="session")
+def nortek_thermostat_state_fixture():
+    """Load the nortek thermostat node state fixture data."""
+    return json.loads(load_fixture("zwave_js/nortek_thermostat_state.json"))
+
+
 @pytest.fixture(name="client")
 def mock_client_fixture(controller_state, version_state):
     """Mock a client."""
@@ -140,6 +147,32 @@ def climate_radio_thermostat_ct100_plus_fixture(
     return node
 
 
+@pytest.fixture(name="nortek_thermostat")
+def nortek_thermostat_fixture(client, nortek_thermostat_state):
+    """Mock a nortek thermostat node."""
+    node = Node(client, nortek_thermostat_state)
+    client.driver.controller.nodes[node.node_id] = node
+    return node
+
+
+@pytest.fixture(name="nortek_thermostat_added_event")
+def nortek_thermostat_added_event_fixture(client):
+    """Mock a Nortek thermostat node added event."""
+    event_data = json.loads(load_fixture("zwave_js/nortek_thermostat_added_event.json"))
+    event = Event("node added", event_data)
+    return event
+
+
+@pytest.fixture(name="nortek_thermostat_removed_event")
+def nortek_thermostat_removed_event_fixture(client):
+    """Mock a Nortek thermostat node removed event."""
+    event_data = json.loads(
+        load_fixture("zwave_js/nortek_thermostat_removed_event.json")
+    )
+    event = Event("node removed", event_data)
+    return event
+
+
 @pytest.fixture(name="integration")
 async def integration_fixture(hass, client):
     """Set up the zwave_js integration."""