Allow Z-Wave controller migration on USB discovery (#143677)

Allow migration on USB discovery

homeassistant/components/zwave_js/config_flow.py

@@ -428,10 +428,19 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
         """Handle USB Discovery."""
         if not is_hassio(self.hass):
             return self.async_abort(reason="discovery_requires_supervisor")
-        if self._async_current_entries():
-            return self.async_abort(reason="already_configured")
         if self._async_in_progress():
             return self.async_abort(reason="already_in_progress")
+        if current_config_entries := self._async_current_entries(include_ignore=False):
+            config_entry = next(
+                (
+                    entry
+                    for entry in current_config_entries
+                    if entry.data.get(CONF_USE_ADDON)
+                ),
+                None,
+            )
+            if not config_entry:
+                return self.async_abort(reason="addon_required")
 
         vid = discovery_info.vid
         pid = discovery_info.pid
@@ -443,7 +452,10 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
             return self.async_abort(reason="not_zwave_device")
 
         addon_info = await self._async_get_addon_info()
-        if addon_info.state not in (AddonState.NOT_INSTALLED, AddonState.NOT_RUNNING):
+        if (
+            addon_info.state not in (AddonState.NOT_INSTALLED, AddonState.INSTALLING)
+            and addon_info.options.get(CONF_ADDON_DEVICE) == discovery_info.device
+        ):
             return self.async_abort(reason="already_configured")
 
         await self.async_set_unique_id(
@@ -482,6 +494,18 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
             )
 
         self._usb_discovery = True
+        if current_config_entries := self._async_current_entries(include_ignore=False):
+            self._reconfigure_config_entry = next(
+                (
+                    entry
+                    for entry in current_config_entries
+                    if entry.data.get(CONF_USE_ADDON)
+                ),
+                None,
+            )
+            if not self._reconfigure_config_entry:
+                return self.async_abort(reason="addon_required")
+            return await self.async_step_intent_migrate()
 
         return await self.async_step_on_supervisor({CONF_USE_ADDON: True})
 
@@ -840,6 +864,14 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
         """Reset the current controller, and instruct the user to unplug it."""
 
         if user_input is not None:
+            config_entry = self._reconfigure_config_entry
+            assert config_entry is not None
+            # Unload the config entry before stopping the add-on.
+            await self.hass.config_entries.async_unload(config_entry.entry_id)
+            if self.usb_path:
+                # USB discovery was used, so the device is already known.
+                await self._async_set_addon_config({CONF_ADDON_DEVICE: self.usb_path})
+                return await self.async_step_start_addon()
             # Now that the old controller is gone, we can scan for serial ports again
             return await self.async_step_choose_serial_port()
 

tests/components/zwave_js/test_config_flow.py

@@ -849,6 +849,134 @@ async def test_usb_discovery_addon_not_running(
     assert len(mock_setup_entry.mock_calls) == 1
 
 
+@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(
+    hass: HomeAssistant,
+    addon_options: dict[str, Any],
+    set_addon_options: AsyncMock,
+    restart_addon: AsyncMock,
+    client: MagicMock,
+    integration: MockConfigEntry,
+) -> None:
+    """Test usb discovery migration."""
+    addon_options["device"] = "/dev/ttyUSB0"
+    entry = integration
+    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.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"
+
+    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"
+
+    await hass.async_block_till_done()
+    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
+
+
 async def test_discovery_addon_not_running(
     hass: HomeAssistant,
     supervisor,
@@ -1072,10 +1200,10 @@ async def test_abort_usb_discovery_with_existing_flow(
     assert result2["reason"] == "already_in_progress"
 
 
-async def test_abort_usb_discovery_already_configured(
+async def test_abort_usb_discovery_addon_required(
     hass: HomeAssistant, supervisor, addon_options
 ) -> None:
-    """Test usb discovery flow is aborted when there is an existing entry."""
+    """Test usb discovery aborted when existing entry not using add-on."""
     entry = MockConfigEntry(
         domain=DOMAIN,
         data={"url": "ws://localhost:3000"},
@@ -1090,7 +1218,52 @@ async def test_abort_usb_discovery_already_configured(
         data=USB_DISCOVERY_INFO,
     )
     assert result["type"] is FlowResultType.ABORT
-    assert result["reason"] == "already_configured"
+    assert result["reason"] == "addon_required"
+
+
+@pytest.mark.usefixtures(
+    "supervisor",
+    "addon_running",
+)
+async def test_abort_usb_discovery_confirm_addon_required(
+    hass: HomeAssistant,
+    addon_options: dict[str, Any],
+) -> None:
+    """Test usb discovery confirm aborted when existing entry not using add-on."""
+    addon_options["device"] = "/dev/another_device"
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            "url": "ws://localhost:3000",
+            "usb_path": "/dev/another_device",
+            "use_addon": True,
+        },
+        title=TITLE,
+        unique_id="1234",
+    )
+    entry.add_to_hass(hass)
+
+    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"
+
+    hass.config_entries.async_update_entry(
+        entry,
+        data={
+            **entry.data,
+            "use_addon": False,
+        },
+    )
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "addon_required"
 
 
 async def test_usb_discovery_requires_supervisor(hass: HomeAssistant) -> None:
@@ -1104,10 +1277,13 @@ async def test_usb_discovery_requires_supervisor(hass: HomeAssistant) -> None:
     assert result["reason"] == "discovery_requires_supervisor"
 
 
-async def test_usb_discovery_already_running(
-    hass: HomeAssistant, supervisor, addon_running
+@pytest.mark.usefixtures("supervisor", "addon_running")
+async def test_usb_discovery_same_device(
+    hass: HomeAssistant,
+    addon_options: dict[str, Any],
 ) -> None:
-    """Test usb discovery flow is aborted when the addon is running."""
+    """Test usb discovery flow is aborted when the add-on device is discovered."""
+    addon_options["device"] = USB_DISCOVERY_INFO.device
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": config_entries.SOURCE_USB},
@@ -3326,8 +3502,6 @@ async def test_reconfigure_migrate_with_addon(
 
     client.driver.controller.async_restore_nvm = AsyncMock(side_effect=mock_restore_nvm)
 
-    hass.config_entries.async_reload = AsyncMock()
-
     events = async_capture_events(
         hass, data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESS_UPDATE
     )
@@ -3375,6 +3549,7 @@ async def test_reconfigure_migrate_with_addon(
         },
     )
 
+    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(
@@ -3391,7 +3566,7 @@ async def test_reconfigure_migrate_with_addon(
     assert result["step_id"] == "restore_nvm"
 
     await hass.async_block_till_done()
-    assert hass.config_entries.async_reload.called
+    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