From 40e3038775dc35d2d39cf61a0f25661deeed1f66 Mon Sep 17 00:00:00 2001 From: Martin Hjelmare Date: Tue, 6 May 2025 15:26:45 +0200 Subject: [PATCH] Fix Z-Wave to reload config entry after migration nvm restore (#144338) --- .../components/zwave_js/config_flow.py | 17 +- tests/components/zwave_js/test_config_flow.py | 301 ++++++++++++++++++ 2 files changed, 317 insertions(+), 1 deletion(-) diff --git a/homeassistant/components/zwave_js/config_flow.py b/homeassistant/components/zwave_js/config_flow.py index c6624046a00..46d9e061f0b 100644 --- a/homeassistant/components/zwave_js/config_flow.py +++ b/homeassistant/components/zwave_js/config_flow.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from contextlib import suppress from datetime import datetime import logging from pathlib import Path @@ -77,6 +78,7 @@ 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 = { @@ -1317,15 +1319,28 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): event["bytesWritten"] / event["total"] * 0.5 + 0.5 ) - controller = self._get_driver().controller + @callback + def set_driver_ready(event: dict) -> None: + "Set the driver ready event." + wait_driver_ready.set() + + driver = self._get_driver() + controller = driver.controller + wait_driver_ready = asyncio.Event() unsubs = [ controller.on("nvm convert progress", forward_progress), controller.on("nvm restore progress", forward_progress), + driver.once("driver ready", set_driver_ready), ] try: await controller.async_restore_nvm(self.backup_data) except FailedCommand as err: raise AbortFlow(f"Failed to restore network: {err}") from err + else: + with suppress(TimeoutError): + async with asyncio.timeout(RESTORE_NVM_DRIVER_READY_TIMEOUT): + await wait_driver_ready.wait() + await self.hass.config_entries.async_reload(config_entry.entry_id) finally: for unsub in unsubs: unsub() diff --git a/tests/components/zwave_js/test_config_flow.py b/tests/components/zwave_js/test_config_flow.py index 08f0ffad4bd..de76d9d9dc4 100644 --- a/tests/components/zwave_js/test_config_flow.py +++ b/tests/components/zwave_js/test_config_flow.py @@ -190,6 +190,19 @@ def mock_sdk_version(client: MagicMock) -> Generator[None]: client.driver.controller.data["sdkVersion"] = original_sdk_version +@pytest.fixture(name="driver_ready_timeout") +def mock_driver_ready_timeout() -> Generator[None]: + """Mock migration nvm restore driver ready timeout.""" + with patch( + ( + "homeassistant.components.zwave_js.config_flow." + "RESTORE_NVM_DRIVER_READY_TIMEOUT" + ), + new=0, + ): + yield + + async def test_manual(hass: HomeAssistant) -> None: """Test we create an entry with manual step.""" @@ -889,6 +902,144 @@ async def test_usb_discovery_migration( """Test usb discovery migration.""" addon_options["device"] = "/dev/ttyUSB0" entry = integration + assert client.connect.call_count == 1 + hass.config_entries.async_update_entry( + entry, + unique_id="1234", + data={ + "url": "ws://localhost:3000", + "use_addon": True, + "usb_path": "/dev/ttyUSB0", + }, + ) + + async def mock_backup_nvm_raw(): + await asyncio.sleep(0) + client.driver.controller.emit( + "nvm backup progress", {"bytesRead": 100, "total": 200} + ) + return b"test_nvm_data" + + client.driver.controller.async_backup_nvm_raw = AsyncMock( + side_effect=mock_backup_nvm_raw + ) + + async def mock_restore_nvm(data: bytes): + client.driver.controller.emit( + "nvm convert progress", + {"event": "nvm convert progress", "bytesRead": 100, "total": 200}, + ) + await asyncio.sleep(0) + client.driver.controller.emit( + "nvm restore progress", + {"event": "nvm restore progress", "bytesWritten": 100, "total": 200}, + ) + client.driver.emit( + "driver ready", {"event": "driver ready", "source": "driver"} + ) + + client.driver.controller.async_restore_nvm = AsyncMock(side_effect=mock_restore_nvm) + + events = async_capture_events( + hass, data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESS_UPDATE + ) + + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_USB}, + data=USB_DISCOVERY_INFO, + ) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "usb_confirm" + assert mock_usb_serial_by_id.call_count == 2 + + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + + assert result["type"] == FlowResultType.FORM + assert result["step_id"] == "intent_migrate" + + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + + assert result["type"] == FlowResultType.SHOW_PROGRESS + assert result["step_id"] == "backup_nvm" + + with patch("pathlib.Path.write_bytes", MagicMock()) as mock_file: + await hass.async_block_till_done() + assert client.driver.controller.async_backup_nvm_raw.call_count == 1 + assert mock_file.call_count == 1 + assert len(events) == 1 + assert events[0].data["progress"] == 0.5 + events.clear() + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert result["type"] == FlowResultType.FORM + assert result["step_id"] == "instruct_unplug" + + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + + assert entry.state is config_entries.ConfigEntryState.NOT_LOADED + assert result["type"] == FlowResultType.SHOW_PROGRESS + assert result["step_id"] == "start_addon" + assert set_addon_options.call_args == call( + "core_zwave_js", AddonsOptions(config={"device": USB_DISCOVERY_INFO.device}) + ) + + await hass.async_block_till_done() + + assert restart_addon.call_args == call("core_zwave_js") + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert result["type"] == FlowResultType.SHOW_PROGRESS + assert result["step_id"] == "restore_nvm" + assert client.connect.call_count == 2 + + await hass.async_block_till_done() + assert client.connect.call_count == 3 + assert entry.state is config_entries.ConfigEntryState.LOADED + assert client.driver.controller.async_restore_nvm.call_count == 1 + assert len(events) == 2 + assert events[0].data["progress"] == 0.25 + assert events[1].data["progress"] == 0.75 + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert result["type"] is FlowResultType.ABORT + assert result["reason"] == "migration_successful" + assert integration.data["url"] == "ws://host1:3001" + assert integration.data["usb_path"] == USB_DISCOVERY_INFO.device + assert integration.data["use_addon"] is True + + +@pytest.mark.usefixtures("supervisor", "addon_running", "get_addon_discovery_info") +@pytest.mark.parametrize( + "discovery_info", + [ + [ + Discovery( + addon="core_zwave_js", + service="zwave_js", + uuid=uuid4(), + config=ADDON_DISCOVERY_INFO, + ) + ] + ], +) +async def test_usb_discovery_migration_driver_ready_timeout( + hass: HomeAssistant, + addon_options: dict[str, Any], + driver_ready_timeout: None, + mock_usb_serial_by_id: MagicMock, + set_addon_options: AsyncMock, + restart_addon: AsyncMock, + client: MagicMock, + integration: MockConfigEntry, +) -> None: + """Test driver ready timeout after nvm restore during usb discovery migration.""" + addon_options["device"] = "/dev/ttyUSB0" + entry = integration + assert client.connect.call_count == 1 hass.config_entries.async_update_entry( entry, unique_id="1234", @@ -976,8 +1127,10 @@ async def test_usb_discovery_migration( assert result["type"] == FlowResultType.SHOW_PROGRESS assert result["step_id"] == "restore_nvm" + assert client.connect.call_count == 2 await hass.async_block_till_done() + assert client.connect.call_count == 3 assert entry.state is config_entries.ConfigEntryState.LOADED assert client.driver.controller.async_restore_nvm.call_count == 1 assert len(events) == 2 @@ -3552,6 +3705,152 @@ async def test_reconfigure_migrate_with_addon( ) -> None: """Test migration flow with add-on.""" entry = integration + assert client.connect.call_count == 1 + hass.config_entries.async_update_entry( + entry, + unique_id="1234", + data={ + "url": "ws://localhost:3000", + "use_addon": True, + "usb_path": "/dev/ttyUSB0", + }, + ) + + async def mock_backup_nvm_raw(): + await asyncio.sleep(0) + client.driver.controller.emit( + "nvm backup progress", {"bytesRead": 100, "total": 200} + ) + return b"test_nvm_data" + + client.driver.controller.async_backup_nvm_raw = AsyncMock( + side_effect=mock_backup_nvm_raw + ) + + async def mock_restore_nvm(data: bytes): + client.driver.controller.emit( + "nvm convert progress", + {"event": "nvm convert progress", "bytesRead": 100, "total": 200}, + ) + await asyncio.sleep(0) + client.driver.controller.emit( + "nvm restore progress", + {"event": "nvm restore progress", "bytesWritten": 100, "total": 200}, + ) + client.driver.emit( + "driver ready", {"event": "driver ready", "source": "driver"} + ) + + client.driver.controller.async_restore_nvm = AsyncMock(side_effect=mock_restore_nvm) + + events = async_capture_events( + hass, data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESS_UPDATE + ) + + result = await entry.start_reconfigure_flow(hass) + + assert result["type"] == FlowResultType.MENU + assert result["step_id"] == "reconfigure" + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {"next_step_id": "intent_migrate"} + ) + + assert result["type"] == FlowResultType.FORM + assert result["step_id"] == "intent_migrate" + + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + + assert result["type"] == FlowResultType.SHOW_PROGRESS + assert result["step_id"] == "backup_nvm" + + with patch("pathlib.Path.write_bytes", MagicMock()) as mock_file: + await hass.async_block_till_done() + assert client.driver.controller.async_backup_nvm_raw.call_count == 1 + assert mock_file.call_count == 1 + assert len(events) == 1 + assert events[0].data["progress"] == 0.5 + events.clear() + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert result["type"] == FlowResultType.FORM + assert result["step_id"] == "instruct_unplug" + + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + + assert result["type"] == FlowResultType.FORM + assert result["step_id"] == "choose_serial_port" + assert result["data_schema"].schema[CONF_USB_PATH] + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ + CONF_USB_PATH: "/test", + }, + ) + + assert entry.state is config_entries.ConfigEntryState.NOT_LOADED + assert result["type"] == FlowResultType.SHOW_PROGRESS + assert result["step_id"] == "start_addon" + assert set_addon_options.call_args == call( + "core_zwave_js", AddonsOptions(config={"device": "/test"}) + ) + + await hass.async_block_till_done() + + assert restart_addon.call_args == call("core_zwave_js") + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert result["type"] == FlowResultType.SHOW_PROGRESS + assert result["step_id"] == "restore_nvm" + assert client.connect.call_count == 2 + + await hass.async_block_till_done() + assert client.connect.call_count == 3 + assert entry.state is config_entries.ConfigEntryState.LOADED + assert client.driver.controller.async_restore_nvm.call_count == 1 + assert len(events) == 2 + assert events[0].data["progress"] == 0.25 + assert events[1].data["progress"] == 0.75 + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert result["type"] is FlowResultType.ABORT + assert result["reason"] == "migration_successful" + assert integration.data["url"] == "ws://host1:3001" + assert integration.data["usb_path"] == "/test" + assert integration.data["use_addon"] is True + + +@pytest.mark.parametrize( + "discovery_info", + [ + [ + Discovery( + addon="core_zwave_js", + service="zwave_js", + uuid=uuid4(), + config=ADDON_DISCOVERY_INFO, + ) + ] + ], +) +async def test_reconfigure_migrate_driver_ready_timeout( + hass: HomeAssistant, + client, + supervisor, + integration, + addon_running, + driver_ready_timeout: None, + restart_addon, + set_addon_options, + get_addon_discovery_info, +) -> None: + """Test migration flow with driver ready timeout after nvm restore.""" + entry = integration + assert client.connect.call_count == 1 hass.config_entries.async_update_entry( entry, unique_id="1234", @@ -3648,8 +3947,10 @@ async def test_reconfigure_migrate_with_addon( assert result["type"] == FlowResultType.SHOW_PROGRESS assert result["step_id"] == "restore_nvm" + assert client.connect.call_count == 2 await hass.async_block_till_done() + assert client.connect.call_count == 3 assert entry.state is config_entries.ConfigEntryState.LOADED assert client.driver.controller.async_restore_nvm.call_count == 1 assert len(events) == 2