Fix Z-Wave restore nvm command to wait for driver ready (#144413)

homeassistant/components/zwave_js/api.py

@@ -88,6 +88,7 @@ from .const import (
     DATA_CLIENT,
     DOMAIN,
     EVENT_DEVICE_ADDED_TO_REGISTRY,
+    RESTORE_NVM_DRIVER_READY_TIMEOUT,
     USER_AGENT,
 )
 from .helpers import (
@@ -3063,14 +3064,28 @@ async def websocket_restore_nvm(
             )
         )
 
+    @callback
+    def set_driver_ready(event: dict) -> None:
+        "Set the driver ready event."
+        wait_driver_ready.set()
+
+    wait_driver_ready = asyncio.Event()
+
     # Set up subscription for progress events
     connection.subscriptions[msg["id"]] = async_cleanup
     msg[DATA_UNSUBSCRIBE] = unsubs = [
         controller.on("nvm convert progress", forward_progress),
         controller.on("nvm restore progress", forward_progress),
+        driver.once("driver ready", set_driver_ready),
     ]
 
     await controller.async_restore_nvm_base64(msg["data"])
+
+    with suppress(TimeoutError):
+        async with asyncio.timeout(RESTORE_NVM_DRIVER_READY_TIMEOUT):
+            await wait_driver_ready.wait()
+    await hass.config_entries.async_reload(entry.entry_id)
+
     connection.send_message(
         websocket_api.event_message(
             msg[ID],

homeassistant/components/zwave_js/config_flow.py

@@ -67,6 +67,7 @@ from .const import (
     CONF_USE_ADDON,
     DATA_CLIENT,
     DOMAIN,
+    RESTORE_NVM_DRIVER_READY_TIMEOUT,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -78,7 +79,6 @@ ADDON_SETUP_TIMEOUT = 5
 ADDON_SETUP_TIMEOUT_ROUNDS = 40
 CONF_EMULATE_HARDWARE = "emulate_hardware"
 CONF_LOG_LEVEL = "log_level"
-RESTORE_NVM_DRIVER_READY_TIMEOUT = 60
 SERVER_VERSION_TIMEOUT = 10
 
 ADDON_LOG_LEVELS = {

homeassistant/components/zwave_js/const.py

@@ -201,3 +201,7 @@ COVER_TILT_PROPERTY_KEYS: set[str | int | None] = {
     WindowCoveringPropertyKey.VERTICAL_SLATS_ANGLE,
     WindowCoveringPropertyKey.VERTICAL_SLATS_ANGLE_NO_POSITION,
 }
+
+# Other constants
+
+RESTORE_NVM_DRIVER_READY_TIMEOUT = 60

tests/components/zwave_js/test_api.py

@@ -5518,10 +5518,98 @@ async def test_restore_nvm(
     # Set up mocks for the controller events
     controller = client.driver.controller
 
-    # Test restore success
+    async def async_send_command_driver_ready(
-    with patch.object(
+        message: dict[str, Any],
-        controller, "async_restore_nvm_base64", return_value=None
+        require_schema: int | None = None,
-    ) as mock_restore:
+    ) -> dict:
+        """Send a command and get a response."""
+        client.driver.emit(
+            "driver ready", {"event": "driver ready", "source": "driver"}
+        )
+        return {}
+
+    client.async_send_command.side_effect = async_send_command_driver_ready
+
+    # Send the subscription request
+    await ws_client.send_json_auto_id(
+        {
+            "type": "zwave_js/restore_nvm",
+            "entry_id": integration.entry_id,
+            "data": "dGVzdA==",  # base64 encoded "test"
+        }
+    )
+
+    # Verify the finished event first
+    msg = await ws_client.receive_json()
+    assert msg["type"] == "event"
+    assert msg["event"]["event"] == "finished"
+
+    # Verify subscription success
+    msg = await ws_client.receive_json()
+    assert msg["type"] == "result"
+    assert msg["success"] is True
+
+    # Simulate progress events
+    event = Event(
+        "nvm restore progress",
+        {
+            "source": "controller",
+            "event": "nvm restore progress",
+            "bytesWritten": 25,
+            "total": 100,
+        },
+    )
+    controller.receive_event(event)
+    msg = await ws_client.receive_json()
+    assert msg["event"]["event"] == "nvm restore progress"
+    assert msg["event"]["bytesWritten"] == 25
+    assert msg["event"]["total"] == 100
+
+    event = Event(
+        "nvm restore progress",
+        {
+            "source": "controller",
+            "event": "nvm restore progress",
+            "bytesWritten": 50,
+            "total": 100,
+        },
+    )
+    controller.receive_event(event)
+    msg = await ws_client.receive_json()
+    assert msg["event"]["event"] == "nvm restore progress"
+    assert msg["event"]["bytesWritten"] == 50
+    assert msg["event"]["total"] == 100
+
+    await hass.async_block_till_done()
+
+    # Verify the restore was called
+    # The first call is the relevant one for nvm restore.
+    assert client.async_send_command.call_count == 3
+    assert client.async_send_command.call_args_list[0] == call(
+        {
+            "command": "controller.restore_nvm",
+            "nvmData": "dGVzdA==",
+        },
+        require_schema=14,
+    )
+
+    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.RESTORE_NVM_DRIVER_READY_TIMEOUT",
+        new=0,
+    ):
         # Send the subscription request
         await ws_client.send_json_auto_id(
             {
@@ -5533,6 +5621,7 @@ async def test_restore_nvm(
 
         # Verify the finished event first
         msg = await ws_client.receive_json()
+
         assert msg["type"] == "event"
         assert msg["event"]["event"] == "finished"
 
@@ -5541,48 +5630,25 @@ async def test_restore_nvm(
         assert msg["type"] == "result"
         assert msg["success"] is True
 
-        # Simulate progress events
-        event = Event(
-            "nvm restore progress",
-            {
-                "source": "controller",
-                "event": "nvm restore progress",
-                "bytesWritten": 25,
-                "total": 100,
-            },
-        )
-        controller.receive_event(event)
-        msg = await ws_client.receive_json()
-        assert msg["event"]["event"] == "nvm restore progress"
-        assert msg["event"]["bytesWritten"] == 25
-        assert msg["event"]["total"] == 100
-
-        event = Event(
-            "nvm restore progress",
-            {
-                "source": "controller",
-                "event": "nvm restore progress",
-                "bytesWritten": 50,
-                "total": 100,
-            },
-        )
-        controller.receive_event(event)
-        msg = await ws_client.receive_json()
-        assert msg["event"]["event"] == "nvm restore progress"
-        assert msg["event"]["bytesWritten"] == 50
-        assert msg["event"]["total"] == 100
-
-        # Wait for the restore to complete
         await hass.async_block_till_done()
 
-        # Verify the restore was called
+    # Verify the restore was called
-        assert mock_restore.called
+    # The first call is the relevant one for nvm restore.
+    assert client.async_send_command.call_count == 3
+    assert client.async_send_command.call_args_list[0] == call(
+        {
+            "command": "controller.restore_nvm",
+            "nvmData": "dGVzdA==",
+        },
+        require_schema=14,
+    )
+
+    client.async_send_command.reset_mock()
 
     # Test restore failure
-    with patch.object(
+    with patch(
-        controller,
+        f"{CONTROLLER_PATCH_PREFIX}.async_restore_nvm_base64",
-        "async_restore_nvm_base64",
+        side_effect=FailedZWaveCommand("failed_command", 1, "error message"),
-        side_effect=FailedCommand("failed_command", "Restore failed"),
     ):
         # Send the subscription request
         await ws_client.send_json_auto_id(
@@ -5596,7 +5662,7 @@ async def test_restore_nvm(
         # Verify error response
         msg = await ws_client.receive_json()
         assert not msg["success"]
-        assert msg["error"]["code"] == "Restore failed"
+        assert msg["error"]["code"] == "zwave_error"
 
     # Test entry_id not found
     await ws_client.send_json_auto_id(