From 619fdea5dfeeccd08eff78681913721748623140 Mon Sep 17 00:00:00 2001 From: Martin Hjelmare Date: Wed, 7 May 2025 18:50:45 +0200 Subject: [PATCH] Fix Z-Wave restore nvm command to wait for driver ready (#144413) --- homeassistant/components/zwave_js/api.py | 15 ++ .../components/zwave_js/config_flow.py | 2 +- homeassistant/components/zwave_js/const.py | 4 + tests/components/zwave_js/test_api.py | 152 +++++++++++++----- 4 files changed, 129 insertions(+), 44 deletions(-) diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py index aa2219031d2..f4397737234 100644 --- a/homeassistant/components/zwave_js/api.py +++ b/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], diff --git a/homeassistant/components/zwave_js/config_flow.py b/homeassistant/components/zwave_js/config_flow.py index 84717047fdd..407af9e902b 100644 --- a/homeassistant/components/zwave_js/config_flow.py +++ b/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 = { diff --git a/homeassistant/components/zwave_js/const.py b/homeassistant/components/zwave_js/const.py index 16cf6f748bb..5792fca42a2 100644 --- a/homeassistant/components/zwave_js/const.py +++ b/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 diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py index 2e3d8fd290a..c6ce3d9ac1b 100644 --- a/tests/components/zwave_js/test_api.py +++ b/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 - with patch.object( - controller, "async_restore_nvm_base64", return_value=None - ) as mock_restore: + async def async_send_command_driver_ready( + 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.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 - assert mock_restore.called + # 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 restore failure - with patch.object( - controller, - "async_restore_nvm_base64", - side_effect=FailedCommand("failed_command", "Restore failed"), + with patch( + f"{CONTROLLER_PATCH_PREFIX}.async_restore_nvm_base64", + side_effect=FailedZWaveCommand("failed_command", 1, "error message"), ): # 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(