Improve Z-Wave migration flow (#143673)

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")
 

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",
     [