Add progress to ZHA migration steps (home-assistant#155764)

Co-authored-by: TheJulianJES <[email protected]>

Author: 2 people

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.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 @@ def __init__(self) -> None:
 
         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 @@ async def async_step_migration_strategy_recommended(
         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 @@ async def async_step_maybe_reset_old_radio(
                 # 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 @@ async def async_step_plug_in_old_radio(
         # 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,15 +531,22 @@ async def async_step_skip_reset_old_radio(
         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
     ) -> ConfigFlowResult:
         """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 @@ async def async_step_form_initial_network(
         # 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 @@ async def async_step_choose_automatic_backup(
             ),
         )
 
-    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)},
-                    )
-
-                return await self._async_create_radio_entry()
+        """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",
+            )
 
-            # On rejection, explain why we can't restore
-            return self.async_abort(reason="cannot_restore_backup_no_ieee_confirm")
+        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 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 @@ async def async_step_maybe_reset_old_radio(
 
         # 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)
 

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%]"