From 0aabb112200efc445768a305d9c440b484a46b13 Mon Sep 17 00:00:00 2001 From: Martin Hjelmare Date: Fri, 25 Apr 2025 18:33:19 +0200 Subject: [PATCH] Improve Z-Wave migration flow (#143673) --- .../components/zwave_js/config_flow.py | 43 ++++++++- tests/components/zwave_js/test_config_flow.py | 93 +++++++++++++++++++ 2 files changed, 131 insertions(+), 5 deletions(-) diff --git a/homeassistant/components/zwave_js/config_flow.py b/homeassistant/components/zwave_js/config_flow.py index 2e24ddc070d..cba27daa026 100644 --- a/homeassistant/components/zwave_js/config_flow.py +++ b/homeassistant/components/zwave_js/config_flow.py @@ -201,6 +201,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): self.backup_data: bytes | None = None self.backup_filepath: str | None = None self.use_addon = False + self._migrating = False self._reconfigure_config_entry: ConfigEntry | None = None self._usb_discovery = False @@ -264,6 +265,8 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: """Add-on start failed.""" + if self._migrating: + return self.async_abort(reason="addon_start_failed") if self._reconfigure_config_entry: return await self.async_revert_addon_config(reason="addon_start_failed") return self.async_abort(reason="addon_start_failed") @@ -315,6 +318,8 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): Get add-on discovery info and server version info. Set unique id and abort if already configured. """ + if self._migrating: + return await self.async_step_finish_addon_setup_migrate(user_input) if self._reconfigure_config_entry: return await self.async_step_finish_addon_setup_reconfigure(user_input) return await self.async_step_finish_addon_setup_user(user_input) @@ -774,6 +779,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): return self.async_abort(reason="addon_required") if user_input is not None: + self._migrating = True return await self.async_step_backup_nvm() return self.async_show_form(step_id="intent_migrate") @@ -1072,6 +1078,37 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): """Migration done.""" return self.async_abort(reason="migration_successful") + async def async_step_finish_addon_setup_migrate( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Prepare info needed to complete the config entry update.""" + ws_address = self.ws_address + assert ws_address is not None + version_info = self.version_info + assert version_info is not None + + # We need to wait for the config entry to be reloaded, + # before restoring the backup. + # We will do this in the restore nvm progress task, + # to get a nicer user experience. + self._async_update_entry( + { + "unique_id": str(version_info.home_id), + CONF_URL: ws_address, + CONF_USB_PATH: self.usb_path, + CONF_S0_LEGACY_KEY: self.s0_legacy_key, + CONF_S2_ACCESS_CONTROL_KEY: self.s2_access_control_key, + CONF_S2_AUTHENTICATED_KEY: self.s2_authenticated_key, + CONF_S2_UNAUTHENTICATED_KEY: self.s2_unauthenticated_key, + CONF_LR_S2_ACCESS_CONTROL_KEY: self.lr_s2_access_control_key, + CONF_LR_S2_AUTHENTICATED_KEY: self.lr_s2_authenticated_key, + CONF_USE_ADDON: True, + CONF_INTEGRATION_CREATED_ADDON: self.integration_created_addon, + }, + schedule_reload=False, + ) + return await self.async_step_restore_nvm() + async def async_step_finish_addon_setup_reconfigure( self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: @@ -1100,9 +1137,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): except CannotConnect: return await self.async_revert_addon_config(reason="cannot_connect") - if self.backup_data is None and config_entry.unique_id != str( - self.version_info.home_id - ): + if config_entry.unique_id != str(self.version_info.home_id): return await self.async_revert_addon_config(reason="different_device") self._async_update_entry( @@ -1119,8 +1154,6 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN): CONF_INTEGRATION_CREATED_ADDON: self.integration_created_addon, } ) - if self.backup_data: - return await self.async_step_restore_nvm() return self.async_abort(reason="reconfigure_successful") diff --git a/tests/components/zwave_js/test_config_flow.py b/tests/components/zwave_js/test_config_flow.py index c6b38f39053..d85d3293218 100644 --- a/tests/components/zwave_js/test_config_flow.py +++ b/tests/components/zwave_js/test_config_flow.py @@ -3483,6 +3483,99 @@ async def test_reconfigure_migrate_backup_file_failure( assert result["reason"] == "backup_failed" +@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_reconfigure_migrate_start_addon_failure( + hass: HomeAssistant, + client: MagicMock, + integration: MockConfigEntry, + restart_addon: AsyncMock, + set_addon_options: AsyncMock, +) -> None: + """Test add-on start failure during migration.""" + restart_addon.side_effect = SupervisorError("Boom!") + entry = integration + hass.config_entries.async_update_entry( + entry, unique_id="1234", data={**entry.data, "use_addon": True} + ) + + async def mock_backup_nvm_raw(): + await asyncio.sleep(0) + return b"test_nvm_data" + + client.driver.controller.async_backup_nvm_raw = AsyncMock( + side_effect=mock_backup_nvm_raw + ) + client.driver.controller.async_restore_nvm = AsyncMock( + side_effect=FailedCommand("test_error", "unknown_error") + ) + + 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 + + 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" + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ + CONF_USB_PATH: "/test", + }, + ) + + assert set_addon_options.call_count == 1 + assert set_addon_options.call_args == call( + "core_zwave_js", AddonsOptions(config={"device": "/test"}) + ) + + assert result["type"] == FlowResultType.SHOW_PROGRESS + assert result["step_id"] == "start_addon" + + await hass.async_block_till_done() + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert result["type"] == FlowResultType.ABORT + assert result["reason"] == "addon_start_failed" + + @pytest.mark.parametrize( "discovery_info", [