Fix Z-Wave controller hard reset (#144389)

homeassistant/components/zwave_js/api.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Callable, Coroutine
+from contextlib import suppress
 import dataclasses
 from functools import partial, wraps
 from typing import Any, Concatenate, Literal, cast
@@ -182,6 +184,8 @@ STRATEGY = "strategy"
 # https://github.com/zwave-js/node-zwave-js/blob/master/packages/core/src/security/QR.ts#L41
 MINIMUM_QR_STRING_LENGTH = 52
 
+HARD_RESET_CONTROLLER_DRIVER_READY_TIMEOUT = 60
+
 
 # Helper schemas
 PLANNED_PROVISIONING_ENTRY_SCHEMA = vol.All(
@@ -2816,6 +2820,7 @@ async def websocket_hard_reset_controller(
     driver: Driver,
 ) -> None:
     """Hard reset controller."""
+    unsubs: list[Callable[[], None]]
 
     @callback
     def async_cleanup() -> None:
@@ -2831,13 +2836,28 @@ async def websocket_hard_reset_controller(
             connection.send_result(msg[ID], device.id)
             async_cleanup()
 
+    @callback
+    def set_driver_ready(event: dict) -> None:
+        "Set the driver ready event."
+        wait_driver_ready.set()
+
+    wait_driver_ready = asyncio.Event()
+
     msg[DATA_UNSUBSCRIBE] = unsubs = [
         async_dispatcher_connect(
             hass, EVENT_DEVICE_ADDED_TO_REGISTRY, _handle_device_added
-        )
+        ),
+        driver.once("driver ready", set_driver_ready),
     ]
+
     await driver.async_hard_reset()
 
+    with suppress(TimeoutError):
+        async with asyncio.timeout(HARD_RESET_CONTROLLER_DRIVER_READY_TIMEOUT):
+            await wait_driver_ready.wait()
+
+    await hass.config_entries.async_reload(entry.entry_id)
+
 
 @websocket_api.websocket_command(
     {

tests/components/zwave_js/test_api.py

@@ -5,7 +5,7 @@ from http import HTTPStatus
 from io import BytesIO
 import json
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch
+from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch
 
 import pytest
 from zwave_js_server.const import (
@@ -5078,53 +5078,97 @@ async def test_subscribe_node_statistics(
     assert msg["error"]["code"] == ERR_NOT_LOADED
 
 
-@pytest.mark.skip(
-    reason="The test needs to be updated to reflect what happens when resetting the controller"
-)
 async def test_hard_reset_controller(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
-    client,
+    client: MagicMock,
-    integration,
+    integration: MockConfigEntry,
-    listen_block,
     hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test that the hard_reset_controller WS API call works."""
     entry = integration
     ws_client = await hass_ws_client(hass)
 
-    device = device_registry.async_get_device(
+    async def async_send_command_driver_ready(
-        identifiers={get_device_id(client.driver, client.driver.controller.nodes[1])}
+        message: dict[str, Any],
-    )
+        require_schema: int | None = None,
+    ) -> dict:
+        """Send a command and get a response."""
+        client.driver.emit(
+            "driver ready", {"event": "driver ready", "source": "driver"}
+        )
+        return {}
 
-    client.async_send_command.return_value = {}
+    client.async_send_command.side_effect = async_send_command_driver_ready
-    await ws_client.send_json(
+
+    await ws_client.send_json_auto_id(
         {
-            ID: 1,
             TYPE: "zwave_js/hard_reset_controller",
             ENTRY_ID: entry.entry_id,
         }
     )
-
-    listen_block.set()
-    listen_block.clear()
-    await hass.async_block_till_done()
-
     msg = await ws_client.receive_json()
+
+    device = device_registry.async_get_device(
+        identifiers={get_device_id(client.driver, client.driver.controller.nodes[1])}
+    )
+    assert device is not None
     assert msg["result"] == device.id
     assert msg["success"]
 
-    assert len(client.async_send_command.call_args_list) == 1
+    assert client.async_send_command.call_count == 3
-    assert client.async_send_command.call_args[0][0] == {"command": "driver.hard_reset"}
+    # The first call is the relevant hard reset command.
+    # 25 is the require_schema parameter.
+    assert client.async_send_command.call_args_list[0] == call(
+        {"command": "driver.hard_reset"}, 25
+    )
+
+    client.async_send_command.reset_mock()
+
+    # Test sending command with driver not ready and timeout.
+
+    async def async_send_command_no_driver_ready(
+        message: dict[str, Any],
+        require_schema: int | None = None,
+    ) -> dict:
+        """Send a command and get a response."""
+        return {}
+
+    client.async_send_command.side_effect = async_send_command_no_driver_ready
+
+    with patch(
+        "homeassistant.components.zwave_js.api.HARD_RESET_CONTROLLER_DRIVER_READY_TIMEOUT",
+        new=0,
+    ):
+        await ws_client.send_json_auto_id(
+            {
+                TYPE: "zwave_js/hard_reset_controller",
+                ENTRY_ID: entry.entry_id,
+            }
+        )
+        msg = await ws_client.receive_json()
+
+    device = device_registry.async_get_device(
+        identifiers={get_device_id(client.driver, client.driver.controller.nodes[1])}
+    )
+    assert device is not None
+    assert msg["result"] == device.id
+    assert msg["success"]
+
+    assert client.async_send_command.call_count == 3
+    # The first call is the relevant hard reset command.
+    # 25 is the require_schema parameter.
+    assert client.async_send_command.call_args_list[0] == call(
+        {"command": "driver.hard_reset"}, 25
+    )
 
     # Test FailedZWaveCommand is caught
     with patch(
         "zwave_js_server.model.driver.Driver.async_hard_reset",
         side_effect=FailedZWaveCommand("failed_command", 1, "error message"),
     ):
-        await ws_client.send_json(
+        await ws_client.send_json_auto_id(
             {
-                ID: 2,
                 TYPE: "zwave_js/hard_reset_controller",
                 ENTRY_ID: entry.entry_id,
             }
@@ -5139,9 +5183,8 @@ async def test_hard_reset_controller(
     await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()
 
-    await ws_client.send_json(
+    await ws_client.send_json_auto_id(
         {
-            ID: 3,
             TYPE: "zwave_js/hard_reset_controller",
             ENTRY_ID: entry.entry_id,
         }
@@ -5151,9 +5194,8 @@ async def test_hard_reset_controller(
     assert not msg["success"]
     assert msg["error"]["code"] == ERR_NOT_LOADED
 
-    await ws_client.send_json(
+    await ws_client.send_json_auto_id(
         {
-            ID: 4,
             TYPE: "zwave_js/hard_reset_controller",
             ENTRY_ID: "INVALID",
         }