From ec6d40a51ccedde6d9eab95cb6b5a9056d0d3a69 Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Wed, 5 Nov 2025 05:10:10 -0500 Subject: [PATCH] Add progress to ZHA migration steps (#155764) Co-authored-by: TheJulianJES --- homeassistant/components/zha/config_flow.py | 115 ++++++--- homeassistant/components/zha/strings.json | 53 ++-- tests/components/zha/test_config_flow.py | 253 +++++++++++++++++--- 3 files changed, 333 insertions(+), 88 deletions(-) diff --git a/homeassistant/components/zha/config_flow.py b/homeassistant/components/zha/config_flow.py index ee4dc4e13cf..7364bb49729 100644 --- a/homeassistant/components/zha/config_flow.py +++ b/homeassistant/components/zha/config_flow.py @@ -3,6 +3,7 @@ from __future__ import annotations from abc import abstractmethod +import asyncio import collections from contextlib import suppress from enum import StrEnum @@ -38,7 +39,7 @@ from homeassistant.config_entries import ( ) from homeassistant.const import CONF_NAME from homeassistant.core import HomeAssistant, callback -from homeassistant.data_entry_flow import AbortFlow +from homeassistant.data_entry_flow import AbortFlow, progress_step from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.hassio import is_hassio from homeassistant.helpers.selector import FileSelector, FileSelectorConfig @@ -179,6 +180,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow): """Mixin for common ZHA flow steps and forms.""" _flow_strategy: ZigbeeFlowStrategy | None = None + _overwrite_ieee_during_restore: bool = False _hass: HomeAssistant _title: str @@ -188,6 +190,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow): self._hass = None # type: ignore[assignment] self._radio_mgr = ZhaRadioManager() + self._restore_backup_task: asyncio.Task[None] | None = None @property def hass(self) -> HomeAssistant: @@ -460,6 +463,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow): self._radio_mgr.chosen_backup = self._radio_mgr.backups[0] return await self.async_step_maybe_reset_old_radio() + @progress_step() async def async_step_maybe_reset_old_radio( self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: @@ -493,7 +497,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow): # Old adapter not found or cannot connect, show prompt to plug back in return await self.async_step_plug_in_old_radio() - return await self.async_step_maybe_confirm_ezsp_restore() + return await self.async_step_restore_backup() async def async_step_plug_in_old_radio( self, user_input: dict[str, Any] | None = None @@ -506,7 +510,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow): # Unless the user removes the config entry whilst we try to reset the old radio # for a few seconds and then also unplugs it, we will basically never hit this if not config_entries: - return await self.async_step_maybe_confirm_ezsp_restore() + return await self.async_step_restore_backup() config_entry = config_entries[0] old_device_path = config_entry.data[CONF_DEVICE][CONF_DEVICE_PATH] @@ -527,7 +531,14 @@ class BaseZhaFlow(ConfigEntryBaseFlow): self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: """Skip resetting the old radio and continue with migration.""" - return await self.async_step_maybe_confirm_ezsp_restore() + return await self.async_step_restore_backup() + + async def async_step_pre_plug_in_new_radio( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Strip user_input before showing "plug in new radio" form.""" + # This step is necessary to prevent `user_input` from being passed through + return await self.async_step_plug_in_new_radio() async def async_step_plug_in_new_radio( self, user_input: dict[str, Any] | None = None @@ -535,7 +546,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow): """Prompt user to plug in the new radio if connection fails.""" if user_input is not None: # User confirmed, retry now - return await self.async_step_maybe_confirm_ezsp_restore() + return await self.async_step_restore_backup() assert self._radio_mgr.device_path is not None @@ -606,6 +617,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow): # This step exists only for translations, it does nothing new return await self.async_step_form_new_network(user_input) + @progress_step() async def async_step_form_new_network( self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: @@ -691,53 +703,78 @@ class BaseZhaFlow(ConfigEntryBaseFlow): ), ) - async def async_step_maybe_confirm_ezsp_restore( + async def async_step_restore_backup( self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: - """Confirm restore for EZSP radios that require permanent IEEE writes.""" - if user_input is not None: - if user_input[OVERWRITE_COORDINATOR_IEEE]: - # On confirmation, overwrite destructively - try: - await self._radio_mgr.restore_backup(overwrite_ieee=True) - except HomeAssistantError: - # User unplugged the new adapter, allow retry - return await self.async_step_plug_in_new_radio() - except CannotWriteNetworkSettings as exc: - return self.async_abort( - reason="cannot_restore_backup", - description_placeholders={"error": str(exc)}, - ) + """Restore network backup to new radio.""" + if self._restore_backup_task is None: + self._restore_backup_task = self.hass.async_create_task( + self._radio_mgr.restore_backup( + overwrite_ieee=self._overwrite_ieee_during_restore + ), + "Restore backup", + ) - return await self._async_create_radio_entry() + if not self._restore_backup_task.done(): + return self.async_show_progress( + step_id="restore_backup", + progress_action="restore_backup", + progress_task=self._restore_backup_task, + ) - # On rejection, explain why we can't restore - return self.async_abort(reason="cannot_restore_backup_no_ieee_confirm") - - # On first attempt, just try to restore nondestructively try: - await self._radio_mgr.restore_backup() + await self._restore_backup_task except DestructiveWriteNetworkSettings: - # Restore cannot happen automatically, we need to ask for permission - pass + # If we cannot restore without overwriting the IEEE, ask for confirmation + return self.async_show_progress_done( + next_step_id="pre_confirm_ezsp_ieee_overwrite" + ) except HomeAssistantError: # User unplugged the new adapter, allow retry - return await self.async_step_plug_in_new_radio() + return self.async_show_progress_done(next_step_id="pre_plug_in_new_radio") except CannotWriteNetworkSettings as exc: return self.async_abort( reason="cannot_restore_backup", description_placeholders={"error": str(exc)}, ) - else: - return await self._async_create_radio_entry() + finally: + self._restore_backup_task = None - # If it fails, show the form - return self.async_show_form( - step_id="maybe_confirm_ezsp_restore", - data_schema=vol.Schema( - {vol.Required(OVERWRITE_COORDINATOR_IEEE, default=True): bool} - ), - ) + # Otherwise, proceed to entry creation + return self.async_show_progress_done(next_step_id="create_entry") + + async def async_step_pre_confirm_ezsp_ieee_overwrite( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Strip user_input before showing confirmation form.""" + # This step is necessary to prevent `user_input` from being passed through + return await self.async_step_confirm_ezsp_ieee_overwrite() + + async def async_step_confirm_ezsp_ieee_overwrite( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Show confirmation form for EZSP IEEE address overwrite.""" + if user_input is None: + return self.async_show_form( + step_id="confirm_ezsp_ieee_overwrite", + data_schema=vol.Schema( + {vol.Required(OVERWRITE_COORDINATOR_IEEE, default=True): bool} + ), + ) + + if not user_input[OVERWRITE_COORDINATOR_IEEE]: + return self.async_abort(reason="cannot_restore_backup_no_ieee_confirm") + + self._overwrite_ieee_during_restore = True + return await self.async_step_restore_backup() + + async def async_step_create_entry( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Create the config entry after successful setup/migration.""" + + # This step only exists so that we can create entries from other steps + return await self._async_create_radio_entry() class ZhaConfigFlowHandler(BaseZhaFlow, ConfigFlow, domain=DOMAIN): @@ -1091,7 +1128,7 @@ class ZhaOptionsFlowHandler(BaseZhaFlow, OptionsFlow): # If we are reconfiguring, the old radio will not be available if self._migration_intent is OptionsMigrationIntent.RECONFIGURE: - return await self.async_step_maybe_confirm_ezsp_restore() + return await self.async_step_restore_backup() return await super().async_step_maybe_reset_old_radio(user_input) diff --git a/homeassistant/components/zha/strings.json b/homeassistant/components/zha/strings.json index 9c47efe937d..92ff0a73bed 100644 --- a/homeassistant/components/zha/strings.json +++ b/homeassistant/components/zha/strings.json @@ -16,6 +16,11 @@ "invalid_backup_json": "Invalid backup JSON" }, "flow_title": "{name}", + "progress": { + "form_new_network": "Forming a new Zigbee network.\n\nWe scan for a clear network channel as part of this process, this can take a minute.", + "maybe_reset_old_radio": "Resetting old adapter.\n\nYour old adapter has been backed up and is being factory reset.", + "restore_backup": "Restoring network settings to new adapter.\n\nThis will take a minute." + }, "step": { "choose_automatic_backup": { "data": { @@ -76,8 +81,15 @@ "confirm": { "description": "Do you want to set up {name}?" }, - "confirm_hardware": { - "description": "Do you want to set up {name}?" + "confirm_ezsp_ieee_overwrite": { + "data": { + "overwrite_coordinator_ieee": "Permanently replace the adapter IEEE address" + }, + "description": "Your backup has a different IEEE address than your adapter. For your network to function properly, the IEEE address of your adapter should also be changed.\n\nThis is a permanent operation.", + "title": "Overwrite adapter IEEE address" + }, + "form_new_network": { + "title": "Forming new network" }, "manual_pick_radio_type": { "data": { @@ -100,15 +112,7 @@ "description": "ZHA was not able to automatically detect serial port settings for your adapter. This usually is an issue with the firmware or permissions.\n\nIf you are using firmware with nonstandard settings, enter the serial port settings", "title": "Serial port settings" }, - "maybe_confirm_ezsp_restore": { - "data": { - "overwrite_coordinator_ieee": "Permanently replace the adapter IEEE address" - }, - "description": "Your backup has a different IEEE address than your adapter. For your network to function properly, the IEEE address of your adapter should also be changed.\n\nThis is a permanent operation.", - "title": "Overwrite adapter IEEE address" - }, "maybe_reset_old_radio": { - "description": "A backup was created earlier and your old adapter is being reset as part of the migration.", "title": "Resetting old adapter" }, "plug_in_new_radio": { @@ -127,6 +131,9 @@ }, "title": "Old adapter not found" }, + "restore_backup": { + "title": "Restoring network to new adapter" + }, "upload_manual_backup": { "data": { "uploaded_backup_file": "Upload a file" @@ -1885,6 +1892,11 @@ "invalid_backup_json": "[%key:component::zha::config::error::invalid_backup_json%]" }, "flow_title": "[%key:component::zha::config::flow_title%]", + "progress": { + "form_new_network": "[%key:component::zha::config::progress::form_new_network%]", + "maybe_reset_old_radio": "[%key:component::zha::config::progress::maybe_reset_old_radio%]", + "restore_backup": "[%key:component::zha::config::progress::restore_backup%]" + }, "step": { "choose_automatic_backup": { "data": { @@ -1930,6 +1942,16 @@ "description": "[%key:component::zha::config::step::choose_serial_port::description%]", "title": "[%key:component::zha::config::step::choose_serial_port::title%]" }, + "confirm_ezsp_ieee_overwrite": { + "data": { + "overwrite_coordinator_ieee": "[%key:component::zha::config::step::confirm_ezsp_ieee_overwrite::data::overwrite_coordinator_ieee%]" + }, + "description": "[%key:component::zha::config::step::confirm_ezsp_ieee_overwrite::description%]", + "title": "[%key:component::zha::config::step::confirm_ezsp_ieee_overwrite::title%]" + }, + "form_new_network": { + "title": "[%key:component::zha::config::step::form_new_network::title%]" + }, "init": { "description": "A backup will be performed and ZHA will be stopped. Do you wish to continue?", "title": "Reconfigure ZHA" @@ -1954,12 +1976,8 @@ "description": "[%key:component::zha::config::step::manual_port_config::description%]", "title": "[%key:component::zha::config::step::manual_port_config::title%]" }, - "maybe_confirm_ezsp_restore": { - "data": { - "overwrite_coordinator_ieee": "[%key:component::zha::config::step::maybe_confirm_ezsp_restore::data::overwrite_coordinator_ieee%]" - }, - "description": "[%key:component::zha::config::step::maybe_confirm_ezsp_restore::description%]", - "title": "[%key:component::zha::config::step::maybe_confirm_ezsp_restore::title%]" + "maybe_reset_old_radio": { + "title": "[%key:component::zha::config::step::maybe_reset_old_radio::title%]" }, "plug_in_new_radio": { "description": "[%key:component::zha::config::step::plug_in_new_radio::description%]", @@ -1989,6 +2007,9 @@ }, "title": "Migrate or re-configure" }, + "restore_backup": { + "title": "[%key:component::zha::config::step::restore_backup::title%]" + }, "upload_manual_backup": { "data": { "uploaded_backup_file": "[%key:component::zha::config::step::upload_manual_backup::data::uploaded_backup_file%]" diff --git a/tests/components/zha/test_config_flow.py b/tests/components/zha/test_config_flow.py index a5044ae756d..77b80606361 100644 --- a/tests/components/zha/test_config_flow.py +++ b/tests/components/zha/test_config_flow.py @@ -1,5 +1,6 @@ """Tests for ZHA config flow.""" +import asyncio from collections.abc import Callable, Coroutine, Generator from datetime import timedelta from ipaddress import ip_address @@ -46,6 +47,7 @@ from homeassistant.config_entries import ( SOURCE_USB, SOURCE_USER, SOURCE_ZEROCONF, + ConfigEntriesFlowManager, ConfigEntryState, ConfigFlowResult, ) @@ -186,6 +188,33 @@ def usb_port(device="/dev/ttyUSB1234") -> USBDevice: ) +async def consume_progress_flow( + hass: HomeAssistant, + flow_id: str, + valid_step_ids: tuple[str, ...], + flow_manager: ConfigEntriesFlowManager | None = None, +) -> ConfigFlowResult: + """Consume a progress flow until it is done.""" + if flow_manager is None: + flow_manager = hass.config_entries.flow + + while True: + result = await flow_manager.async_configure(flow_id) + flow_id = result["flow_id"] + + if result["type"] != FlowResultType.SHOW_PROGRESS: + break + + assert result["type"] is FlowResultType.SHOW_PROGRESS + assert result["step_id"] in valid_step_ids + + await asyncio.sleep(0.1) + + # Ensure all background tasks complete + await hass.async_block_till_done() + return result + + @pytest.mark.parametrize( ("entry_name", "unique_id", "radio_type", "service_info"), [ @@ -297,10 +326,16 @@ async def test_zeroconf_discovery( assert result_confirm["type"] is FlowResultType.MENU assert result_confirm["step_id"] == "choose_setup_strategy" - result_form = await hass.config_entries.flow.async_configure( + result_setup = await hass.config_entries.flow.async_configure( result_confirm["flow_id"], user_input={"next_step_id": config_flow.SETUP_STRATEGY_RECOMMENDED}, ) + + result_form = await consume_progress_flow( + hass, + flow_id=result_setup["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result_form["type"] is FlowResultType.CREATE_ENTRY @@ -351,10 +386,16 @@ async def test_legacy_zeroconf_discovery_zigate( assert result_confirm["type"] is FlowResultType.MENU assert result_confirm["step_id"] == "choose_setup_strategy" - result_form = await hass.config_entries.flow.async_configure( + result_setup = await hass.config_entries.flow.async_configure( result_confirm["flow_id"], user_input={"next_step_id": config_flow.SETUP_STRATEGY_RECOMMENDED}, ) + + result_form = await consume_progress_flow( + hass, + flow_id=result_setup["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result_form["type"] is FlowResultType.CREATE_ENTRY @@ -486,10 +527,16 @@ async def test_discovery_via_usb(hass: HomeAssistant) -> None: assert result2["step_id"] == "choose_setup_strategy" with patch("homeassistant.components.zha.async_setup_entry", return_value=True): - result3 = await hass.config_entries.flow.async_configure( + result_setup = await hass.config_entries.flow.async_configure( result2["flow_id"], user_input={"next_step_id": config_flow.SETUP_STRATEGY_RECOMMENDED}, ) + + result3 = await consume_progress_flow( + hass, + flow_id=result_setup["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result3["type"] is FlowResultType.CREATE_ENTRY @@ -618,11 +665,17 @@ async def test_migration_strategy_recommended( return_value=True, ) as mock_async_unload, ): - result_recommended = await hass.config_entries.flow.async_configure( + result_migrate = await hass.config_entries.flow.async_configure( result_confirm["flow_id"], user_input={"next_step_id": config_flow.MIGRATION_STRATEGY_RECOMMENDED}, ) + result_recommended = await consume_progress_flow( + hass, + flow_id=result_migrate["flow_id"], + valid_step_ids=("maybe_reset_old_radio", "restore_backup"), + ) + assert mock_async_unload.mock_calls == [call(entry.entry_id)] assert result_recommended["type"] is FlowResultType.ABORT assert result_recommended["reason"] == "reconfigure_successful" @@ -669,11 +722,17 @@ async def test_migration_strategy_recommended_cannot_write( "homeassistant.components.zha.radio_manager.ZhaRadioManager.restore_backup", side_effect=CannotWriteNetworkSettings("test error"), ) as mock_restore_backup: - result_recommended = await hass.config_entries.flow.async_configure( + result_migrate = await hass.config_entries.flow.async_configure( result_confirm["flow_id"], user_input={"next_step_id": config_flow.MIGRATION_STRATEGY_RECOMMENDED}, ) + result_recommended = await consume_progress_flow( + hass, + flow_id=result_migrate["flow_id"], + valid_step_ids=("maybe_reset_old_radio", "restore_backup"), + ) + assert mock_restore_backup.call_count == 1 assert result_recommended["type"] is FlowResultType.ABORT assert result_recommended["reason"] == "cannot_restore_backup" @@ -1021,10 +1080,16 @@ async def test_user_flow(hass: HomeAssistant) -> None: assert result["step_id"] == "choose_setup_strategy" with patch("homeassistant.components.zha.async_setup_entry", return_value=True): - result2 = await hass.config_entries.flow.async_configure( + result_setup = await hass.config_entries.flow.async_configure( result["flow_id"], user_input={"next_step_id": config_flow.SETUP_STRATEGY_RECOMMENDED}, ) + + result2 = await consume_progress_flow( + hass, + flow_id=result_setup["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result2["type"] is FlowResultType.CREATE_ENTRY @@ -1222,10 +1287,16 @@ async def test_user_port_config(probe_mock, hass: HomeAssistant) -> None: assert result["type"] is FlowResultType.MENU assert result["step_id"] == "choose_setup_strategy" - result2 = await hass.config_entries.flow.async_configure( + result_setup = await hass.config_entries.flow.async_configure( result["flow_id"], user_input={"next_step_id": config_flow.SETUP_STRATEGY_RECOMMENDED}, ) + + result2 = await consume_progress_flow( + hass, + flow_id=result_setup["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert ( @@ -1251,9 +1322,15 @@ async def test_hardware_not_onboarded(hass: HomeAssistant) -> None: with patch( "homeassistant.components.onboarding.async_is_onboarded", return_value=False ): - result_create = await hass.config_entries.flow.async_init( + result_init = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_HARDWARE}, data=data ) + + result_create = await consume_progress_flow( + hass, + flow_id=result_init["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result_create["title"] == "Yellow" @@ -1298,10 +1375,16 @@ async def test_hardware_no_flow_strategy(hass: HomeAssistant) -> None: assert result2["type"] is FlowResultType.MENU assert result2["step_id"] == "choose_setup_strategy" - result_create = await hass.config_entries.flow.async_configure( + result_setup = await hass.config_entries.flow.async_configure( result2["flow_id"], user_input={"next_step_id": config_flow.SETUP_STRATEGY_RECOMMENDED}, ) + + result_create = await consume_progress_flow( + hass, + flow_id=result_setup["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result_create["title"] == "Yellow" @@ -1346,10 +1429,16 @@ async def test_hardware_flow_strategy_advanced(hass: HomeAssistant) -> None: assert confirm_result["type"] is FlowResultType.MENU assert confirm_result["step_id"] == "choose_formation_strategy" - result_create = await hass.config_entries.flow.async_configure( + result_form = await hass.config_entries.flow.async_configure( confirm_result["flow_id"], user_input={"next_step_id": "form_new_network"}, ) + + result_create = await consume_progress_flow( + hass, + flow_id=result_form["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result_create["type"] is FlowResultType.CREATE_ENTRY @@ -1387,10 +1476,16 @@ async def test_hardware_flow_strategy_recommended(hass: HomeAssistant) -> None: assert result_hardware["type"] is FlowResultType.FORM assert result_hardware["step_id"] == "confirm" - result_create = await hass.config_entries.flow.async_configure( + result_confirm = await hass.config_entries.flow.async_configure( result_hardware["flow_id"], user_input={}, ) + + result_create = await consume_progress_flow( + hass, + flow_id=result_confirm["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result_create["type"] is FlowResultType.CREATE_ENTRY @@ -1467,10 +1562,16 @@ async def test_hardware_migration_flow_strategy_advanced( assert result_confirm["type"] is FlowResultType.MENU assert result_confirm["step_id"] == "choose_formation_strategy" - result_formation_strategy = await hass.config_entries.flow.async_configure( + result_form = await hass.config_entries.flow.async_configure( result_confirm["flow_id"], user_input={"next_step_id": "form_new_network"}, ) + + result_formation_strategy = await consume_progress_flow( + hass, + flow_id=result_form["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result_formation_strategy["type"] is FlowResultType.ABORT @@ -1534,10 +1635,16 @@ async def test_hardware_migration_flow_strategy_recommended( assert result_hardware["type"] is FlowResultType.FORM assert result_hardware["step_id"] == "confirm" - result_confirm = await hass.config_entries.flow.async_configure( + result_migrate = await hass.config_entries.flow.async_configure( result_hardware["flow_id"], user_input={} ) + result_confirm = await consume_progress_flow( + hass, + flow_id=result_migrate["flow_id"], + valid_step_ids=("maybe_reset_old_radio", "restore_backup"), + ) + assert result_confirm["type"] is FlowResultType.ABORT assert result_confirm["reason"] == "reconfigure_successful" assert mock_async_unload.mock_calls == [call(entry.entry_id)] @@ -1643,10 +1750,16 @@ async def test_formation_strategy_form_new_network( """Test forming a new network.""" result = await advanced_pick_radio(RadioType.ezsp) - result2 = await hass.config_entries.flow.async_configure( + result_form = await hass.config_entries.flow.async_configure( result["flow_id"], user_input={"next_step_id": config_flow.FORMATION_FORM_NEW_NETWORK}, ) + + result2 = await consume_progress_flow( + hass, + flow_id=result_form["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() # A new network will be formed @@ -1669,10 +1782,16 @@ async def test_formation_strategy_form_initial_network( mock_app.form_network.side_effect = form_network_side_effect result = await advanced_pick_radio(RadioType.ezsp) - result2 = await hass.config_entries.flow.async_configure( + result_form = await hass.config_entries.flow.async_configure( result["flow_id"], user_input={"next_step_id": config_flow.FORMATION_FORM_INITIAL_NETWORK}, ) + + result2 = await consume_progress_flow( + hass, + flow_id=result_form["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() # A new network will be formed @@ -1709,9 +1828,15 @@ async def test_onboarding_auto_formation_new_hardware( with patch( "homeassistant.components.onboarding.async_is_onboarded", return_value=False ): - result = await hass.config_entries.flow.async_init( + result_init = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USB}, data=discovery_info ) + + result = await consume_progress_flow( + hass, + flow_id=result_init["flow_id"], + valid_step_ids=("form_new_network",), + ) await hass.async_block_till_done() assert result["type"] is FlowResultType.CREATE_ENTRY @@ -1781,11 +1906,17 @@ async def test_formation_strategy_restore_manual_backup_non_ezsp( "homeassistant.components.zha.config_flow.ZhaConfigFlowHandler._parse_uploaded_backup", return_value=zigpy.backups.NetworkBackup(), ): - result3 = await hass.config_entries.flow.async_configure( + result_upload = await hass.config_entries.flow.async_configure( result2["flow_id"], user_input={config_flow.UPLOADED_BACKUP_FILE: str(uuid.uuid4())}, ) + result3 = await consume_progress_flow( + hass, + flow_id=result_upload["flow_id"], + valid_step_ids=("restore_backup",), + ) + mock_app.backups.restore_backup.assert_called_once() allow_overwrite_ieee_mock.assert_not_called() @@ -1826,26 +1957,32 @@ async def test_formation_strategy_restore_manual_backup_overwrite_ieee_ezsp( ], ) as mock_restore_backup, ): - result3 = await hass.config_entries.flow.async_configure( + result_upload = await hass.config_entries.flow.async_configure( result2["flow_id"], user_input={config_flow.UPLOADED_BACKUP_FILE: str(uuid.uuid4())}, ) + result3 = await consume_progress_flow( + hass, + flow_id=result_upload["flow_id"], + valid_step_ids=("restore_backup",), + ) + assert mock_restore_backup.call_count == 1 assert not mock_restore_backup.mock_calls[0].kwargs.get("overwrite_ieee") mock_restore_backup.reset_mock() # The radio requires user confirmation for restore assert result3["type"] is FlowResultType.FORM - assert result3["step_id"] == "maybe_confirm_ezsp_restore" + assert result3["step_id"] == "confirm_ezsp_ieee_overwrite" - result4 = await hass.config_entries.flow.async_configure( + result_confirm = await hass.config_entries.flow.async_configure( result3["flow_id"], user_input={config_flow.OVERWRITE_COORDINATOR_IEEE: True}, ) - assert result4["type"] is FlowResultType.CREATE_ENTRY - assert result4["data"][CONF_RADIO_TYPE] == "ezsp" + assert result_confirm["type"] is FlowResultType.CREATE_ENTRY + assert result_confirm["data"][CONF_RADIO_TYPE] == "ezsp" assert mock_restore_backup.call_count == 1 assert mock_restore_backup.mock_calls[0].kwargs["overwrite_ieee"] is True @@ -1883,18 +2020,24 @@ async def test_formation_strategy_restore_manual_backup_ezsp( ], ) as mock_restore_backup, ): - result3 = await hass.config_entries.flow.async_configure( + result_upload = await hass.config_entries.flow.async_configure( result2["flow_id"], user_input={config_flow.UPLOADED_BACKUP_FILE: str(uuid.uuid4())}, ) + result3 = await consume_progress_flow( + hass, + flow_id=result_upload["flow_id"], + valid_step_ids=("restore_backup",), + ) + assert mock_restore_backup.call_count == 1 assert not mock_restore_backup.mock_calls[0].kwargs.get("overwrite_ieee") mock_restore_backup.reset_mock() # The radio requires user confirmation for restore assert result3["type"] is FlowResultType.FORM - assert result3["step_id"] == "maybe_confirm_ezsp_restore" + assert result3["step_id"] == "confirm_ezsp_ieee_overwrite" result4 = await hass.config_entries.flow.async_configure( result3["flow_id"], @@ -1984,13 +2127,19 @@ async def test_formation_strategy_restore_automatic_backup_ezsp( assert result2["type"] is FlowResultType.FORM assert result2["step_id"] == "choose_automatic_backup" - result3 = await hass.config_entries.flow.async_configure( + result_backup = await hass.config_entries.flow.async_configure( result2["flow_id"], user_input={ config_flow.CHOOSE_AUTOMATIC_BACKUP: "choice:" + repr(backup), }, ) + result3 = await consume_progress_flow( + hass, + flow_id=result_backup["flow_id"], + valid_step_ids=("restore_backup",), + ) + mock_app.backups.restore_backup.assert_called_once() assert result3["type"] is FlowResultType.CREATE_ENTRY @@ -2046,13 +2195,19 @@ async def test_formation_strategy_restore_automatic_backup_non_ezsp( f"choice:{mock_app.backups.backups[1]!r}", ] - result3 = await hass.config_entries.flow.async_configure( + result_backup = await hass.config_entries.flow.async_configure( result2["flow_id"], user_input={ config_flow.CHOOSE_AUTOMATIC_BACKUP: f"choice:{backup!r}", }, ) + result3 = await consume_progress_flow( + hass, + flow_id=result_backup["flow_id"], + valid_step_ids=("restore_backup",), + ) + mock_app.backups.restore_backup.assert_called_once_with(backup) assert result3["type"] is FlowResultType.CREATE_ENTRY @@ -2489,13 +2644,20 @@ async def test_options_flow_migration_reset_old_adapter( spec=ZhaRadioManager, side_effect=[mock_radio_manager], ): - result_strategy = await hass.config_entries.options.async_configure( + result_migrate_start = await hass.config_entries.options.async_configure( flow["flow_id"], user_input={ "next_step_id": config_flow.MIGRATION_STRATEGY_RECOMMENDED, }, ) + result_strategy = await consume_progress_flow( + hass, + flow_id=result_migrate_start["flow_id"], + valid_step_ids=("maybe_reset_old_radio", "restore_backup"), + flow_manager=hass.config_entries.options, + ) + # The old adapter is reset, not the new one assert mock_radio_manager.device_path == "/dev/ttyUSB_old" assert mock_radio_manager.async_reset_adapter.call_count == 1 @@ -2582,13 +2744,20 @@ async def test_options_flow_reconfigure_no_reset( with patch( "homeassistant.components.zha.config_flow.ZhaRadioManager" ) as mock_radio_manager: - result_strategy = await hass.config_entries.options.async_configure( + result_migrate_start = await hass.config_entries.options.async_configure( flow["flow_id"], user_input={ "next_step_id": config_flow.MIGRATION_STRATEGY_RECOMMENDED, }, ) + result_strategy = await consume_progress_flow( + hass, + flow_id=result_migrate_start["flow_id"], + valid_step_ids=("maybe_reset_old_radio", "restore_backup"), + flow_manager=hass.config_entries.options, + ) + # A temp radio manager is never created assert mock_radio_manager.call_count == 0 @@ -2821,11 +2990,17 @@ async def test_migration_resets_old_radio( assert result_confirm["step_id"] == "choose_migration_strategy" - result_recommended = await hass.config_entries.flow.async_configure( + result_migrate = await hass.config_entries.flow.async_configure( result_confirm["flow_id"], user_input={"next_step_id": config_flow.MIGRATION_STRATEGY_RECOMMENDED}, ) + result_recommended = await consume_progress_flow( + hass, + flow_id=result_migrate["flow_id"], + valid_step_ids=("maybe_reset_old_radio", "restore_backup"), + ) + assert result_recommended["type"] is FlowResultType.ABORT assert result_recommended["reason"] == "reconfigure_successful" @@ -2906,18 +3081,24 @@ async def test_formation_strategy_restore_manual_backup_overwrite_ieee_ezsp_writ ], ) as mock_restore_backup, ): - confirm_restore_result = await hass.config_entries.flow.async_configure( + result_upload = await hass.config_entries.flow.async_configure( upload_backup_result["flow_id"], user_input={config_flow.UPLOADED_BACKUP_FILE: str(uuid.uuid4())}, ) + confirm_restore_result = await consume_progress_flow( + hass, + flow_id=result_upload["flow_id"], + valid_step_ids=("restore_backup",), + ) + assert mock_restore_backup.call_count == 1 assert not mock_restore_backup.mock_calls[0].kwargs.get("overwrite_ieee") mock_restore_backup.reset_mock() # The radio requires user confirmation for restore assert confirm_restore_result["type"] is FlowResultType.FORM - assert confirm_restore_result["step_id"] == "maybe_confirm_ezsp_restore" + assert confirm_restore_result["step_id"] == "confirm_ezsp_ieee_overwrite" final_result = await hass.config_entries.flow.async_configure( confirm_restore_result["flow_id"], @@ -3040,7 +3221,7 @@ async def test_plug_in_new_radio_retry( # This adapter requires user confirmation for restore assert result4["type"] is FlowResultType.FORM - assert result4["step_id"] == "maybe_confirm_ezsp_restore" + assert result4["step_id"] == "confirm_ezsp_ieee_overwrite" # Confirm destructive rewrite, but adapter is unplugged again result5 = await hass.config_entries.flow.async_configure( @@ -3150,11 +3331,17 @@ async def test_plug_in_old_radio_retry(hass: HomeAssistant, backup, mock_app) -> assert result_retry["step_id"] == "plug_in_old_radio" # Skip resetting the old adapter - result_skip = await hass.config_entries.flow.async_configure( + result_skip_progress = await hass.config_entries.flow.async_configure( result_retry["flow_id"], user_input={"next_step_id": "skip_reset_old_radio"}, ) + result_skip = await consume_progress_flow( + hass, + flow_id=result_skip_progress["flow_id"], + valid_step_ids=("maybe_reset_old_radio", "restore_backup"), + ) + # Entry created successfully after skipping reset assert result_skip["type"] is FlowResultType.ABORT assert result_skip["reason"] == "reconfigure_successful"