Fix zwave_js websocket api KeyError on unloaded entry (#50154)

homeassistant/components/websocket_api/const.py

@@ -23,6 +23,7 @@ MAX_PENDING_MSG = 2048
 ERR_ID_REUSE = "id_reuse"
 ERR_INVALID_FORMAT = "invalid_format"
 ERR_NOT_FOUND = "not_found"
+ERR_NOT_LOADED = "not_loaded"
 ERR_NOT_SUPPORTED = "not_supported"
 ERR_HOME_ASSISTANT_ERROR = "home_assistant_error"
 ERR_UNKNOWN_COMMAND = "unknown_command"

homeassistant/components/zwave_js/api.py

@@ -22,10 +22,11 @@ from homeassistant.components.http.view import HomeAssistantView
 from homeassistant.components.websocket_api.connection import ActiveConnection
 from homeassistant.components.websocket_api.const import (
     ERR_NOT_FOUND,
+    ERR_NOT_LOADED,
     ERR_NOT_SUPPORTED,
     ERR_UNKNOWN_ERROR,
 )
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ENTRY_STATE_LOADED, ConfigEntry
 from homeassistant.const import CONF_URL
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import config_validation as cv
@@ -83,6 +84,13 @@ def async_get_entry(orig_func: Callable) -> Callable:
                 msg[ID], ERR_NOT_FOUND, f"Config entry {entry_id} not found"
             )
             return
+
+        if entry.state != ENTRY_STATE_LOADED:
+            connection.send_error(
+                msg[ID], ERR_NOT_LOADED, f"Config entry {entry_id} not loaded"
+            )
+            return
+
         client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
         await orig_func(hass, connection, msg, entry, client)
 
@@ -137,17 +145,20 @@ def async_register_api(hass: HomeAssistant) -> None:
     hass.http.register_view(DumpView)  # type: ignore
 
 
-@websocket_api.require_admin
+@websocket_api.require_admin  # type: ignore
+@websocket_api.async_response
 @websocket_api.websocket_command(
     {vol.Required(TYPE): "zwave_js/network_status", vol.Required(ENTRY_ID): str}
 )
-@callback
+@async_get_entry
-def websocket_network_status(
+async def websocket_network_status(
-    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+    hass: HomeAssistant,
+    connection: ActiveConnection,
+    msg: dict,
+    entry: ConfigEntry,
+    client: Client,
 ) -> None:
     """Get the status of the Z-Wave JS network."""
-    entry_id = msg[ENTRY_ID]
-    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
     data = {
         "client": {
             "ws_server_url": client.ws_server_url,
@@ -166,6 +177,7 @@ def websocket_network_status(
     )
 
 
+@websocket_api.async_response  # type: ignore
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zwave_js/node_status",
@@ -173,20 +185,14 @@ def websocket_network_status(
         vol.Required(NODE_ID): int,
     }
 )
-@callback
+@async_get_node
-def websocket_node_status(
+async def websocket_node_status(
-    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+    hass: HomeAssistant,
+    connection: ActiveConnection,
+    msg: dict,
+    node: Node,
 ) -> None:
     """Get the status of a Z-Wave JS node."""
-    entry_id = msg[ENTRY_ID]
-    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
-    node_id = msg[NODE_ID]
-    node = client.driver.controller.nodes.get(node_id)
-
-    if node is None:
-        connection.send_error(msg[ID], ERR_NOT_FOUND, f"Node {node_id} not found")
-        return
-
     data = {
         "node_id": node.node_id,
         "is_routing": node.is_routing,
@@ -537,7 +543,8 @@ async def websocket_set_config_parameter(
     )
 
 
-@websocket_api.require_admin
+@websocket_api.require_admin  # type: ignore
+@websocket_api.async_response
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zwave_js/get_config_parameters",
@@ -545,20 +552,11 @@ async def websocket_set_config_parameter(
         vol.Required(NODE_ID): int,
     }
 )
-@callback
+@async_get_node
-def websocket_get_config_parameters(
+async def websocket_get_config_parameters(
-    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict, node: Node
 ) -> None:
     """Get a list of configuration parameters for a Z-Wave node."""
-    entry_id = msg[ENTRY_ID]
-    node_id = msg[NODE_ID]
-    client = hass.data[DOMAIN][entry_id][DATA_CLIENT]
-    node = client.driver.controller.nodes.get(node_id)
-
-    if node is None:
-        connection.send_error(msg[ID], ERR_NOT_FOUND, f"Node {node_id} not found")
-        return
-
     values = node.get_configuration_values()
     result = {}
     for value_id, zwave_value in values.items():

tests/components/zwave_js/test_api.py

@@ -6,7 +6,7 @@ from zwave_js_server.const import LogLevel
 from zwave_js_server.event import Event
 from zwave_js_server.exceptions import InvalidNewValue, NotFoundError, SetValueFailed
 
-from homeassistant.components.websocket_api.const import ERR_NOT_FOUND
+from homeassistant.components.websocket_api.const import ERR_NOT_FOUND, ERR_NOT_LOADED
 from homeassistant.components.zwave_js.api import (
     COMMAND_CLASS_ID,
     CONFIG,
@@ -31,8 +31,8 @@ from homeassistant.components.zwave_js.const import (
 from homeassistant.helpers import device_registry as dr
 
 
-async def test_websocket_api(hass, integration, multisensor_6, hass_ws_client):
+async def test_network_status(hass, integration, hass_ws_client):
-    """Test the network and node status websocket commands."""
+    """Test the network status websocket command."""
     entry = integration
     ws_client = await hass_ws_client(hass)
 
@@ -45,6 +45,24 @@ async def test_websocket_api(hass, integration, multisensor_6, hass_ws_client):
     assert result["client"]["ws_server_url"] == "ws://test:3000/zjs"
     assert result["client"]["server_version"] == "1.0.0"
 
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {ID: 3, TYPE: "zwave_js/network_status", ENTRY_ID: entry.entry_id}
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
+
+async def test_node_status(hass, integration, multisensor_6, hass_ws_client):
+    """Test the node status websocket command."""
+    entry = integration
+    ws_client = await hass_ws_client(hass)
+
     node = multisensor_6
     await ws_client.send_json(
         {
@@ -63,33 +81,10 @@ async def test_websocket_api(hass, integration, multisensor_6, hass_ws_client):
     assert not result["is_secure"]
     assert result["status"] == 1
 
-    # Test getting configuration parameter values
-    await ws_client.send_json(
-        {
-            ID: 4,
-            TYPE: "zwave_js/get_config_parameters",
-            ENTRY_ID: entry.entry_id,
-            NODE_ID: node.node_id,
-        }
-    )
-    msg = await ws_client.receive_json()
-    result = msg["result"]
-
-    assert len(result) == 61
-    key = "52-112-0-2"
-    assert result[key]["property"] == 2
-    assert result[key]["property_key"] is None
-    assert result[key]["metadata"]["type"] == "number"
-    assert result[key]["configuration_value_type"] == "enumerated"
-    assert result[key]["metadata"]["states"]
-
-    key = "52-112-0-201-255"
-    assert result[key]["property_key"] == 255
-
     # Test getting non-existent node fails
     await ws_client.send_json(
         {
-            ID: 5,
+            ID: 4,
             TYPE: "zwave_js/node_status",
             ENTRY_ID: entry.entry_id,
             NODE_ID: 99999,
@@ -99,18 +94,22 @@ async def test_websocket_api(hass, integration, multisensor_6, hass_ws_client):
     assert not msg["success"]
     assert msg["error"]["code"] == ERR_NOT_FOUND
 
-    # Test getting non-existent node config params fails
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
     await ws_client.send_json(
         {
-            ID: 6,
+            ID: 5,
-            TYPE: "zwave_js/get_config_parameters",
+            TYPE: "zwave_js/node_status",
             ENTRY_ID: entry.entry_id,
-            NODE_ID: 99999,
+            NODE_ID: node.node_id,
         }
     )
     msg = await ws_client.receive_json()
+
     assert not msg["success"]
-    assert msg["error"]["code"] == ERR_NOT_FOUND
+    assert msg["error"]["code"] == ERR_NOT_LOADED
 
 
 async def test_add_node(
@@ -145,6 +144,29 @@ async def test_add_node(
     client.driver.receive_event(nortek_thermostat_added_event)
     msg = await ws_client.receive_json()
     assert msg["event"]["event"] == "node added"
+    node_details = {
+        "node_id": 53,
+        "status": 0,
+        "ready": False,
+    }
+    assert msg["event"]["node"] == node_details
+
+    msg = await ws_client.receive_json()
+    assert msg["event"]["event"] == "device registered"
+    # Check the keys of the device item
+    assert list(msg["event"]["device"]) == ["name", "id"]
+
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {ID: 4, TYPE: "zwave_js/add_node", ENTRY_ID: entry.entry_id}
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
 
 
 async def test_cancel_inclusion_exclusion(hass, integration, client, hass_ws_client):
@@ -168,6 +190,26 @@ async def test_cancel_inclusion_exclusion(hass, integration, client, hass_ws_cli
     msg = await ws_client.receive_json()
     assert msg["success"]
 
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {ID: 6, TYPE: "zwave_js/stop_inclusion", ENTRY_ID: entry.entry_id}
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
+    await ws_client.send_json(
+        {ID: 7, TYPE: "zwave_js/stop_exclusion", ENTRY_ID: entry.entry_id}
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_remove_node(
     hass,
@@ -226,6 +268,18 @@ async def test_remove_node(
     )
     assert device is None
 
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {ID: 4, TYPE: "zwave_js/remove_node", ENTRY_ID: entry.entry_id}
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_refresh_node_info(
     hass, client, integration, hass_ws_client, multisensor_6
@@ -295,6 +349,36 @@ async def test_refresh_node_info(
 
     client.async_send_command_no_wait.reset_mock()
 
+    # Test getting non-existent node fails
+    await ws_client.send_json(
+        {
+            ID: 2,
+            TYPE: "zwave_js/refresh_node_info",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: 9999,
+        }
+    )
+    msg = await ws_client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_FOUND
+
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {
+            ID: 3,
+            TYPE: "zwave_js/refresh_node_info",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: 52,
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_refresh_node_values(
     hass, client, integration, hass_ws_client, multisensor_6
@@ -391,6 +475,38 @@ async def test_refresh_node_cc_values(
     assert not msg["success"]
     assert msg["error"]["code"] == ERR_NOT_FOUND
 
+    # Test getting non-existent node fails
+    await ws_client.send_json(
+        {
+            ID: 3,
+            TYPE: "zwave_js/refresh_node_cc_values",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: 9999,
+            COMMAND_CLASS_ID: 112,
+        }
+    )
+    msg = await ws_client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_FOUND
+
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {
+            ID: 4,
+            TYPE: "zwave_js/refresh_node_cc_values",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: 52,
+            COMMAND_CLASS_ID: 112,
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_set_config_parameter(
     hass, client, hass_ws_client, multisensor_6, integration
@@ -510,6 +626,103 @@ async def test_set_config_parameter(
         assert msg["error"]["code"] == "unknown_error"
         assert msg["error"]["message"] == "test"
 
+    # Test getting non-existent node fails
+    await ws_client.send_json(
+        {
+            ID: 5,
+            TYPE: "zwave_js/set_config_parameter",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: 9999,
+            PROPERTY: 102,
+            PROPERTY_KEY: 1,
+            VALUE: 1,
+        }
+    )
+    msg = await ws_client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_FOUND
+
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {
+            ID: 6,
+            TYPE: "zwave_js/set_config_parameter",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: 52,
+            PROPERTY: 102,
+            PROPERTY_KEY: 1,
+            VALUE: 1,
+        }
+    )
+
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
+
+async def test_get_config_parameters(hass, integration, multisensor_6, hass_ws_client):
+    """Test the get config parameters websocket command."""
+    entry = integration
+    ws_client = await hass_ws_client(hass)
+    node = multisensor_6
+
+    # Test getting configuration parameter values
+    await ws_client.send_json(
+        {
+            ID: 4,
+            TYPE: "zwave_js/get_config_parameters",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: node.node_id,
+        }
+    )
+    msg = await ws_client.receive_json()
+    result = msg["result"]
+
+    assert len(result) == 61
+    key = "52-112-0-2"
+    assert result[key]["property"] == 2
+    assert result[key]["property_key"] is None
+    assert result[key]["metadata"]["type"] == "number"
+    assert result[key]["configuration_value_type"] == "enumerated"
+    assert result[key]["metadata"]["states"]
+
+    key = "52-112-0-201-255"
+    assert result[key]["property_key"] == 255
+
+    # Test getting non-existent node config params fails
+    await ws_client.send_json(
+        {
+            ID: 5,
+            TYPE: "zwave_js/get_config_parameters",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: 99999,
+        }
+    )
+    msg = await ws_client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_FOUND
+
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {
+            ID: 6,
+            TYPE: "zwave_js/get_config_parameters",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: node.node_id,
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_dump_view(integration, hass_client):
     """Test the HTTP dump view."""
@@ -571,6 +784,18 @@ async def test_subscribe_logs(hass, integration, client, hass_ws_client):
         "timestamp": "time",
     }
 
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {ID: 2, TYPE: "zwave_js/subscribe_logs", ENTRY_ID: entry.entry_id}
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_update_log_config(hass, client, integration, hass_ws_client):
     """Test that the update_log_config WS API call works and that schema validation works."""
@@ -691,6 +916,23 @@ async def test_update_log_config(hass, client, integration, hass_ws_client):
         and "must be provided if logging to file" in msg["error"]["message"]
     )
 
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {
+            ID: 7,
+            TYPE: "zwave_js/update_log_config",
+            ENTRY_ID: entry.entry_id,
+            CONFIG: {LEVEL: "Error"},
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_get_log_config(hass, client, integration, hass_ws_client):
     """Test that the get_log_config WS API call works."""
@@ -726,6 +968,22 @@ async def test_get_log_config(hass, client, integration, hass_ws_client):
     assert log_config["filename"] == "/test.txt"
     assert log_config["force_console"] is False
 
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {
+            ID: 2,
+            TYPE: "zwave_js/get_log_config",
+            ENTRY_ID: entry.entry_id,
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
 
 async def test_data_collection(hass, client, integration, hass_ws_client):
     """Test that the data collection WS API commands work."""
@@ -794,3 +1052,32 @@ async def test_data_collection(hass, client, integration, hass_ws_client):
     assert not entry.data[CONF_DATA_COLLECTION_OPTED_IN]
 
     client.async_send_command.reset_mock()
+
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json(
+        {
+            ID: 4,
+            TYPE: "zwave_js/data_collection_status",
+            ENTRY_ID: entry.entry_id,
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
+    await ws_client.send_json(
+        {
+            ID: 5,
+            TYPE: "zwave_js/update_data_collection_preference",
+            ENTRY_ID: entry.entry_id,
+            OPTED_IN: True,
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED