From e14a356c24b4b154dfa5295ca7e71944ad79b28f Mon Sep 17 00:00:00 2001 From: Martin Hjelmare Date: Sat, 26 Apr 2025 07:52:32 +0200 Subject: [PATCH] Allow Z-Wave controller migration on USB discovery (#143677) Allow migration on USB discovery --- .../components/zwave_js/config_flow.py | 38 +++- tests/components/zwave_js/test_config_flow.py | 193 +++++++++++++++++- 2 files changed, 219 insertions(+), 12 deletions(-) diff --git a/homeassistant/components/zwave_js/config_flow.py b/homeassistant/components/zwave_js/config_flow.py index cba27daa026..b453764aa4e 100644 --- a/homeassistant/components/zwave_js/config_flow.py +++ b/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() diff --git a/tests/components/zwave_js/test_config_flow.py b/tests/components/zwave_js/test_config_flow.py index d85d3293218..c5ccd615f5c 100644 --- a/tests/components/zwave_js/test_config_flow.py +++ b/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