From a52b73e6c05c2ef07a620b60a605fb2b7e2a0847 Mon Sep 17 00:00:00 2001 From: Petar Petrov Date: Thu, 27 Aug 2026 15:59:26 +0300 Subject: [PATCH] Add Z-Wave JS WebSocket command to get network neighbors (#179363) --- homeassistant/components/zwave_js/__init__.py | 14 +- homeassistant/components/zwave_js/api.py | 83 +++ homeassistant/components/zwave_js/models.py | 3 + tests/components/zwave_js/test_api.py | 473 ++++++++++++++++++ tests/components/zwave_js/test_init.py | 34 +- 5 files changed, 605 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/zwave_js/__init__.py b/homeassistant/components/zwave_js/__init__.py index 98039f5e3d9c..97befc644491 100644 --- a/homeassistant/components/zwave_js/__init__.py +++ b/homeassistant/components/zwave_js/__init__.py @@ -235,9 +235,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ZwaveJSConfigEntry) -> b entry.async_on_unload(client.disconnect) + # Local because runtime_data is not set yet if HA shuts down during setup + network_neighbors_lock = asyncio.Lock() + async def handle_ha_shutdown(event: Event) -> None: """Handle HA shutdown.""" - await client.disconnect() + # Wait for a running network neighbors refresh, so the client is not + # disconnected before it has turned the radio back on + async with network_neighbors_lock: + await client.disconnect() entry.async_on_unload( hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, handle_ha_shutdown) @@ -270,6 +276,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ZwaveJSConfigEntry) -> b entry_runtime_data = ZwaveJSData( client=client, driver_events=driver_events, + network_neighbors_lock=network_neighbors_lock, ) entry.runtime_data = entry_runtime_data @@ -1167,6 +1174,11 @@ async def client_listen( async def async_unload_entry(hass: HomeAssistant, entry: ZwaveJSConfigEntry) -> bool: """Unload a config entry.""" + # Wait for a running network neighbors refresh, so the client is not + # disconnected before it has turned the radio back on + async with entry.runtime_data.network_neighbors_lock: + pass + unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) entry_runtime_data = entry.runtime_data diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py index 5b98bdfd69e9..6e100af630cd 100644 --- a/homeassistant/components/zwave_js/api.py +++ b/homeassistant/components/zwave_js/api.py @@ -1,9 +1,11 @@ """Websocket API for Z-Wave JS.""" +import asyncio from collections.abc import Callable, Coroutine from contextlib import suppress import dataclasses from functools import partial, wraps +import logging from typing import TYPE_CHECKING, Any, Concatenate, Literal, cast from aiohttp import web, web_exceptions, web_request @@ -104,12 +106,15 @@ if TYPE_CHECKING: from .models import ZwaveJSConfigEntry +_LOGGER = logging.getLogger(__name__) + DATA_UNSUBSCRIBE = "unsubs" # general API constants ID = "id" ENTRY_ID = "entry_id" ERR_NOT_LOADED = "not_loaded" +ERR_RF_TOGGLE_FAILED = "rf_toggle_failed" NODE_ID = "node_id" DEVICE_ID = "device_id" COMMAND_CLASS_ID = "command_class_id" @@ -412,6 +417,7 @@ def async_register_api(hass: HomeAssistant) -> None: websocket_api.async_register_command(hass, websocket_network_status) websocket_api.async_register_command(hass, websocket_subscribe_node_status) websocket_api.async_register_command(hass, websocket_node_status) + websocket_api.async_register_command(hass, websocket_network_neighbors) websocket_api.async_register_command(hass, websocket_node_metadata) websocket_api.async_register_command(hass, websocket_node_alerts) websocket_api.async_register_command(hass, websocket_add_node) @@ -618,6 +624,83 @@ async def websocket_node_status( connection.send_result(msg[ID], node_status(node)) +@websocket_api.require_admin +@websocket_api.websocket_command( + { + vol.Required(TYPE): "zwave_js/network_neighbors", + vol.Required(ENTRY_ID): str, + } +) +@websocket_api.async_response +@async_handle_failed_command +@async_get_entry +async def websocket_network_neighbors( + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict[str, Any], + entry: ZwaveJSConfigEntry, + client: Client, + driver: Driver, +) -> None: + """Get the node IDs of the neighbors of all nodes in the network. + + Reading the routing table can wedge older controllers when the radio is + on or reads overlap, so refreshes are serialized and done with RF off: + https://zwave-js.github.io/zwave-js/#/api/controller?id=getnodeneighbors + """ + controller = driver.controller + + async def restore_rf() -> bool: + """Turn the radio back on, returning False instead of raising.""" + try: + return await controller.async_toggle_rf(True) + except BaseZwaveJSServerError: + return False + + async def read_network_neighbors() -> tuple[bool, bool, dict[int, list[int]]]: + """Read the neighbors of all nodes while the radio is off.""" + neighbors: dict[int, list[int]] = {} + rf_disabled = False + async with entry.runtime_data.network_neighbors_lock: + try: + rf_disabled = await controller.async_toggle_rf(False) + if rf_disabled: + # Snapshot the nodes, inclusion/exclusion can mutate the + # collection while it is being iterated + for node in list(controller.nodes.values()): + # Long range nodes are not part of the mesh + if node.protocol is Protocols.ZWAVE_LONG_RANGE: + continue + try: + neighbors[ + node.node_id + ] = await controller.async_get_node_neighbors(node) + except FailedCommand: + continue + finally: + rf_restored = await restore_rf() + if not rf_restored: + _LOGGER.error( + "Failed to re-enable RF after reading the neighbors of" + " the nodes of config entry %s", + entry.entry_id, + ) + return rf_disabled, rf_restored, neighbors + + # The refresh runs as its own task and is only abandoned on cancellation, + # so a closing connection can't interrupt it while the radio is off + rf_disabled, rf_restored, neighbors = await asyncio.shield( + hass.async_create_task(read_network_neighbors()) + ) + if not rf_disabled: + connection.send_error(msg[ID], ERR_RF_TOGGLE_FAILED, "Failed to disable RF") + return + if not rf_restored: + connection.send_error(msg[ID], ERR_RF_TOGGLE_FAILED, "Failed to re-enable RF") + return + connection.send_result(msg[ID], neighbors) + + @websocket_api.websocket_command( { vol.Required(TYPE): "zwave_js/node_metadata", diff --git a/homeassistant/components/zwave_js/models.py b/homeassistant/components/zwave_js/models.py index 24bc1ced79a0..b42972e221bf 100644 --- a/homeassistant/components/zwave_js/models.py +++ b/homeassistant/components/zwave_js/models.py @@ -1,5 +1,6 @@ """Provide models for the Z-Wave integration.""" +import asyncio from collections.abc import Iterable from dataclasses import asdict, dataclass, field from enum import StrEnum @@ -32,6 +33,8 @@ class ZwaveJSData: client: ZwaveClient driver_events: DriverEvents old_server_log_level: LogLevel | None = None + # Serializes routing table reads, which require the radio to be off + network_neighbors_lock: asyncio.Lock = field(default_factory=asyncio.Lock) type ZwaveJSConfigEntry = ConfigEntry[ZwaveJSData] diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py index b417c244b39c..742c23316049 100644 --- a/tests/components/zwave_js/test_api.py +++ b/tests/components/zwave_js/test_api.py @@ -1,6 +1,7 @@ """Test the Z-Wave JS Websocket API.""" import asyncio +from collections.abc import Callable, Coroutine from copy import deepcopy from http import HTTPStatus from io import BytesIO @@ -405,6 +406,478 @@ async def test_node_status( assert msg["error"]["code"] == ERR_NOT_LOADED +def mock_neighbors_commands( + client: MagicMock, + handler: Callable[[dict[str, Any]], Coroutine[Any, Any, dict[str, Any]]], +) -> list[dict[str, Any]]: + """Record the commands sent to the driver and answer them with handler.""" + commands: list[dict[str, Any]] = [] + + async def _send_command(message: dict[str, Any], **kwargs: Any) -> dict[str, Any]: + commands.append(message) + return await handler(message) + + client.async_send_command.side_effect = _send_command + return commands + + +async def neighbors_ok(message: dict[str, Any]) -> dict[str, Any]: + """Answer both toggling RF and reading neighbors successfully.""" + if message["command"] == "controller.get_node_neighbors": + return {"neighbors": []} + return {"success": True} + + +async def test_network_neighbors( + hass: HomeAssistant, + multisensor_6: Node, + wallmote_central_scene: Node, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test the network_neighbors websocket command.""" + ws_client = await hass_ws_client(hass) + # Long range nodes are not part of the mesh and must be skipped + wallmote_central_scene.data["protocol"] = Protocols.ZWAVE_LONG_RANGE + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.get_node_neighbors": + neighbors = [35, 32] if message["nodeId"] == multisensor_6.node_id else [] + return {"neighbors": neighbors} + return {"success": True} + + commands = mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + msg = await ws_client.receive_json() + + assert msg["success"] + assert msg["result"] == { + "1": [], + str(multisensor_6.node_id): [35, 32], + } + # The nodes are read one at a time while the radio is off + assert commands == [ + {"command": "controller.toggle_rf", "enabled": False}, + {"command": "controller.get_node_neighbors", "nodeId": 1}, + {"command": "controller.get_node_neighbors", "nodeId": multisensor_6.node_id}, + {"command": "controller.toggle_rf", "enabled": True}, + ] + + +async def test_network_neighbors_node_failure( + hass: HomeAssistant, + multisensor_6: Node, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test a node that fails to report its neighbors is skipped.""" + ws_client = await hass_ws_client(hass) + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] != "controller.get_node_neighbors": + return {"success": True} + if message["nodeId"] == 1: + raise FailedZWaveCommand("failed_command", 1, "error message") + return {"neighbors": []} + + commands = mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + msg = await ws_client.receive_json() + + assert msg["success"] + assert msg["result"] == {str(multisensor_6.node_id): []} + assert commands[-1] == {"command": "controller.toggle_rf", "enabled": True} + + +async def test_network_neighbors_node_added_while_reading( + hass: HomeAssistant, + multisensor_6: Node, + wallmote_central_scene: Node, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test a node joining the network while reading doesn't abort the reads.""" + ws_client = await hass_ws_client(hass) + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.get_node_neighbors": + client.driver.controller.nodes.setdefault(999, wallmote_central_scene) + return {"neighbors": []} + return {"success": True} + + commands = mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + msg = await ws_client.receive_json() + + assert msg["success"] + assert msg["result"] == { + "1": [], + str(multisensor_6.node_id): [], + str(wallmote_central_scene.node_id): [], + } + assert commands[-1] == {"command": "controller.toggle_rf", "enabled": True} + + +async def test_network_neighbors_rf_disable_rejected( + hass: HomeAssistant, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test the nodes are not read when the radio can't be turned off.""" + ws_client = await hass_ws_client(hass) + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.toggle_rf": + return {"success": message["enabled"]} + return {"neighbors": []} + + commands = mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + msg = await ws_client.receive_json() + + assert not msg["success"] + assert msg["error"]["code"] == "rf_toggle_failed" + assert commands == [ + {"command": "controller.toggle_rf", "enabled": False}, + {"command": "controller.toggle_rf", "enabled": True}, + ] + + +async def test_network_neighbors_rf_restore_rejected( + hass: HomeAssistant, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test a radio that can't be turned back on is reported and logged.""" + ws_client = await hass_ws_client(hass) + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.toggle_rf": + return {"success": not message["enabled"]} + return {"neighbors": []} + + commands = mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + msg = await ws_client.receive_json() + + assert not msg["success"] + assert msg["error"]["code"] == "rf_toggle_failed" + assert commands[-1] == {"command": "controller.toggle_rf", "enabled": True} + assert "Failed to re-enable RF" in caplog.text + + +async def test_network_neighbors_rf_toggle_error( + hass: HomeAssistant, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test a single response is sent when the radio can't be toggled at all.""" + ws_client = await hass_ws_client(hass) + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["enabled"]: + raise FailedZWaveCommand("failed_command", 1, "error message") + return {"success": False} + + mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + msg = await ws_client.receive_json() + + assert not msg["success"] + assert msg["error"]["code"] == "rf_toggle_failed" + assert "Failed to re-enable RF" in caplog.text + + # The next frame is the pong, proving no second response was sent + await ws_client.send_json_auto_id({TYPE: "ping"}) + msg = await ws_client.receive_json() + assert msg["type"] == "pong" + + +async def test_network_neighbors_cancelled( + hass: HomeAssistant, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test the radio is turned back on when the command is cancelled.""" + ws_client = await hass_ws_client(hass) + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.toggle_rf" and not message["enabled"]: + raise asyncio.CancelledError + return {"success": True} + + commands = mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + # A cancelled command sends no response, so sync on a ping instead + await ws_client.send_json_auto_id({TYPE: "ping"}) + msg = await ws_client.receive_json() + assert msg["type"] == "pong" + await hass.async_block_till_done() + + assert commands == [ + {"command": "controller.toggle_rf", "enabled": False}, + {"command": "controller.toggle_rf", "enabled": True}, + ] + + +async def test_network_neighbors_handler_cancelled( + hass: HomeAssistant, + multisensor_6: Node, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test an abandoned request doesn't interrupt the refresh. + + The refresh must keep the lock and turn the radio back on even when the + websocket command handler is cancelled, e.g. by a closing connection. + """ + ws_client = await hass_ws_client(hass) + read_started = asyncio.Event() + resume_read = asyncio.Event() + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.get_node_neighbors": + read_started.set() + await resume_read.wait() + return {"neighbors": []} + return {"success": True} + + commands = mock_neighbors_commands(client, handler) + lock = integration.runtime_data.network_neighbors_lock + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + await read_started.wait() + + handler_task = next( + task + for task in asyncio.all_tasks() + if "_handle_async_response" in repr(task.get_coro()) + ) + handler_task.cancel() + for _ in range(5): + await asyncio.sleep(0) + # The abandoned refresh keeps reading with the lock held + assert lock.locked() + + resume_read.set() + await hass.async_block_till_done() + assert not lock.locked() + assert commands == [ + {"command": "controller.toggle_rf", "enabled": False}, + {"command": "controller.get_node_neighbors", "nodeId": 1}, + {"command": "controller.get_node_neighbors", "nodeId": multisensor_6.node_id}, + {"command": "controller.toggle_rf", "enabled": True}, + ] + + +async def test_network_neighbors_concurrent( + hass: HomeAssistant, + multisensor_6: Node, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test concurrent requests are serialized so the radio is never shared.""" + ws_client = await hass_ws_client(hass) + ws_client_2 = await hass_ws_client(hass) + read_started = asyncio.Event() + resume_read = asyncio.Event() + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.get_node_neighbors": + read_started.set() + await resume_read.wait() + return {"neighbors": []} + return {"success": True} + + commands = mock_neighbors_commands(client, handler) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + await read_started.wait() + + await ws_client_2.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + await ws_client_2.send_json_auto_id({TYPE: "ping"}) + msg = await ws_client_2.receive_json() + assert msg["type"] == "pong" + # The second request must not have touched the radio yet + assert commands == [ + {"command": "controller.toggle_rf", "enabled": False}, + {"command": "controller.get_node_neighbors", "nodeId": 1}, + ] + + resume_read.set() + msg = await ws_client.receive_json() + assert msg["success"] + msg = await ws_client_2.receive_json() + assert msg["success"] + # The second request turns the radio off only after the first turned it on + assert commands == [ + {"command": "controller.toggle_rf", "enabled": False}, + {"command": "controller.get_node_neighbors", "nodeId": 1}, + {"command": "controller.get_node_neighbors", "nodeId": multisensor_6.node_id}, + {"command": "controller.toggle_rf", "enabled": True}, + {"command": "controller.toggle_rf", "enabled": False}, + {"command": "controller.get_node_neighbors", "nodeId": 1}, + {"command": "controller.get_node_neighbors", "nodeId": multisensor_6.node_id}, + {"command": "controller.toggle_rf", "enabled": True}, + ] + + +async def test_network_neighbors_unload_waits( + hass: HomeAssistant, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test unloading the entry waits for the radio to be turned back on.""" + ws_client = await hass_ws_client(hass) + read_started = asyncio.Event() + resume_read = asyncio.Event() + + async def handler(message: dict[str, Any]) -> dict[str, Any]: + if message["command"] == "controller.get_node_neighbors": + read_started.set() + await resume_read.wait() + return {"neighbors": []} + return {"success": True} + + commands = mock_neighbors_commands(client, handler) + + async def mock_disconnect() -> None: + commands.append({"command": "disconnect"}) + + client.disconnect.side_effect = mock_disconnect + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + await read_started.wait() + + unload_task = hass.async_create_task( + hass.config_entries.async_unload(integration.entry_id) + ) + for _ in range(10): + await asyncio.sleep(0) + # The refresh is holding the lock, so the client must still be connected + assert not unload_task.done() + + resume_read.set() + msg = await ws_client.receive_json() + assert msg["success"] + assert await unload_task + # The radio was turned back on before the client disconnected + assert commands[-2:] == [ + {"command": "controller.toggle_rf", "enabled": True}, + {"command": "disconnect"}, + ] + + +async def test_network_neighbors_invalid_entry( + hass: HomeAssistant, + integration: MockConfigEntry, + client: MagicMock, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test the network_neighbors websocket command with an invalid entry.""" + ws_client = await hass_ws_client(hass) + mock_neighbors_commands(client, neighbors_ok) + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: "fake_entry_id", + } + ) + msg = await ws_client.receive_json() + + assert not msg["success"] + assert msg["error"]["code"] == ERR_NOT_FOUND + + await hass.config_entries.async_unload(integration.entry_id) + await hass.async_block_till_done() + + await ws_client.send_json_auto_id( + { + TYPE: "zwave_js/network_neighbors", + ENTRY_ID: integration.entry_id, + } + ) + msg = await ws_client.receive_json() + + assert not msg["success"] + assert msg["error"]["code"] == ERR_NOT_LOADED + + async def test_node_metadata( hass: HomeAssistant, wallmote_central_scene, diff --git a/tests/components/zwave_js/test_init.py b/tests/components/zwave_js/test_init.py index ba533229dd01..d4cb60165c0d 100644 --- a/tests/components/zwave_js/test_init.py +++ b/tests/components/zwave_js/test_init.py @@ -26,7 +26,12 @@ from homeassistant.components.persistent_notification import async_dismiss from homeassistant.components.zwave_js import DOMAIN from homeassistant.components.zwave_js.helpers import get_device_id, get_device_id_ext from homeassistant.config_entries import ConfigEntryDisabler, ConfigEntryState -from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN, Platform +from homeassistant.const import ( + EVENT_HOMEASSISTANT_STOP, + STATE_UNAVAILABLE, + STATE_UNKNOWN, + Platform, +) from homeassistant.core import CoreState, HomeAssistant from homeassistant.helpers import ( area_registry as ar, @@ -158,6 +163,33 @@ async def test_home_assistant_stop( assert client.disconnect.call_count == 1 +async def test_home_assistant_stop_waits_for_neighbors_refresh( + hass: HomeAssistant, + integration: MockConfigEntry, + client: MagicMock, +) -> None: + """Test stop disconnects only after a neighbors refresh releases the lock. + + The refresh turns the radio back on before releasing, so disconnecting + earlier could leave the radio off. + """ + disconnected = asyncio.Event() + + async def mock_disconnect() -> None: + disconnected.set() + + client.disconnect.side_effect = mock_disconnect + + async with integration.runtime_data.network_neighbors_lock: + hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) + for _ in range(10): + await asyncio.sleep(0) + assert not disconnected.is_set() + + await hass.async_block_till_done() + assert disconnected.is_set() + + @pytest.mark.usefixtures("client", "connect_timeout") async def test_initialized_timeout(hass: HomeAssistant) -> None: """Test we handle a timeout during client initialization."""