Add ZHA migration retry steps for unplugged adapters

homeassistant/components/zha/config_flow.py

@@ -473,10 +473,62 @@ async def async_step_maybe_reset_old_radio(
             temp_radio_mgr.device_settings = config_entry.data[CONF_DEVICE]
             temp_radio_mgr.radio_type = RadioType[config_entry.data[CONF_RADIO_TYPE]]
 
-            await temp_radio_mgr.async_reset_adapter()
+            try:
+                await temp_radio_mgr.async_reset_adapter()
+            except HomeAssistantError:
+                # 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()
 
+    async def async_step_plug_in_old_radio(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Prompt user to plug in the old radio if connection fails."""
+        config_entries = self.hass.config_entries.async_entries(
+            DOMAIN, include_ignore=False
+        )
+
+        # config entry should always exist here, skip if not
+        if not config_entries:
+            return await self.async_step_maybe_confirm_ezsp_restore()
+
+        config_entry = config_entries[0]
+        old_device_path = config_entry.data[CONF_DEVICE][CONF_DEVICE_PATH]
+
+        return self.async_show_menu(
+            step_id="plug_in_old_radio",
+            menu_options=["retry_old_radio", "skip_reset_old_radio"],
+            description_placeholders={"device_path": old_device_path},
+        )
+
+    async def async_step_retry_old_radio(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Retry connecting to the old radio."""
+        return await self.async_step_maybe_reset_old_radio()
+
+    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()
+
+    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()
+
+        assert self._radio_mgr.device_path is not None
+
+        return self.async_show_form(
+            step_id="plug_in_new_radio",
+            description_placeholders={"device_path": self._radio_mgr.device_path},
+        )
+
     async def async_step_migration_strategy_advanced(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -633,6 +685,9 @@ async def async_step_maybe_confirm_ezsp_restore(
                 # 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",
@@ -650,6 +705,9 @@ async def async_step_maybe_confirm_ezsp_restore(
         except DestructiveWriteNetworkSettings:
             # Restore cannot happen automatically, we need to ask for permission
             pass
+        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",

homeassistant/components/zha/strings.json

@@ -111,6 +111,22 @@
         "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": {
+        "description": "Your new adapter at `{device_path}` was not found.\nPlease plug it in and click Submit to continue.",
+        "title": "New adapter not found"
+      },
+      "plug_in_old_radio": {
+        "description": "Your old adapter at `{device_path}` was not found. You can retry after plugging it back in, or skip resetting the old adapter.\n\nIf you skip resetting the old adapter, please make sure it is permanently unavailable, as plugging it back in later will cause network issues.",
+        "menu_option_descriptions": {
+          "retry_old_radio": "Retry connecting to the old adapter to reset it as part of the migration.",
+          "skip_reset_old_radio": "Skip resetting the old adapter and continue with the migration."
+        },
+        "menu_options": {
+          "retry_old_radio": "Retry",
+          "skip_reset_old_radio": "Skip reset"
+        },
+        "title": "Old adapter not found"
+      },
       "upload_manual_backup": {
         "data": {
           "uploaded_backup_file": "Upload a file"
@@ -1945,6 +1961,22 @@
         "description": "[%key:component::zha::config::step::maybe_confirm_ezsp_restore::description%]",
         "title": "[%key:component::zha::config::step::maybe_confirm_ezsp_restore::title%]"
       },
+      "plug_in_new_radio": {
+        "description": "[%key:component::zha::config::step::plug_in_new_radio::description%]",
+        "title": "[%key:component::zha::config::step::plug_in_new_radio::title%]"
+      },
+      "plug_in_old_radio": {
+        "description": "[%key:component::zha::config::step::plug_in_old_radio::description%]",
+        "menu_option_descriptions": {
+          "retry_old_radio": "[%key:component::zha::config::step::plug_in_old_radio::menu_option_descriptions::retry_old_radio%]",
+          "skip_reset_old_radio": "[%key:component::zha::config::step::plug_in_old_radio::menu_option_descriptions::skip_reset_old_radio%]"
+        },
+        "menu_options": {
+          "retry_old_radio": "[%key:component::zha::config::step::plug_in_old_radio::menu_options::retry_old_radio%]",
+          "skip_reset_old_radio": "[%key:component::zha::config::step::plug_in_old_radio::menu_options::skip_reset_old_radio%]"
+        },
+        "title": "[%key:component::zha::config::step::plug_in_old_radio::title%]"
+      },
       "prompt_migrate_or_reconfigure": {
         "description": "Are you migrating to a new adapter or re-configuring the current adapter?",
         "menu_option_descriptions": {

tests/components/zha/test_config_flow.py

@@ -51,6 +51,7 @@
 from homeassistant.const import CONF_SOURCE
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.service_info.ssdp import (
     ATTR_UPNP_MANUFACTURER_URL,
     ATTR_UPNP_SERIAL,
@@ -2925,3 +2926,183 @@ async def test_migrate_setup_options_with_ignored_discovery(
         "setup_strategy_recommended",
         "setup_strategy_advanced",
     ]
+
+
+@patch("homeassistant.components.zha.radio_manager._allow_overwrite_ezsp_ieee")
+async def test_plug_in_new_radio_retry(
+    allow_overwrite_ieee_mock,
+    advanced_pick_radio: RadioPicker,
+    mock_app: AsyncMock,
+    backup,
+    hass: HomeAssistant,
+) -> None:
+    """Test plug_in_new_radio step when restore fails due to unplugged adapter."""
+    result = await advanced_pick_radio(RadioType.ezsp)
+
+    result2 = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        user_input={"next_step_id": config_flow.FORMATION_UPLOAD_MANUAL_BACKUP},
+    )
+    await hass.async_block_till_done()
+
+    assert result2["type"] is FlowResultType.FORM
+    assert result2["step_id"] == "upload_manual_backup"
+
+    with (
+        patch(
+            "homeassistant.components.zha.config_flow.ZhaConfigFlowHandler._parse_uploaded_backup",
+            return_value=backup,
+        ),
+        patch(
+            "homeassistant.components.zha.radio_manager.ZhaRadioManager.restore_backup",
+            side_effect=[
+                HomeAssistantError(
+                    "Failed to connect to Zigbee adapter: [Errno 2] No such file or directory"
+                ),
+                DestructiveWriteNetworkSettings("Radio IEEE change is permanent"),
+                HomeAssistantError(
+                    "Failed to connect to Zigbee adapter: [Errno 2] No such file or directory"
+                ),
+                None,
+            ],
+        ) as mock_restore_backup,
+    ):
+        result3 = await hass.config_entries.flow.async_configure(
+            result2["flow_id"],
+            user_input={config_flow.UPLOADED_BACKUP_FILE: str(uuid.uuid4())},
+        )
+
+        # Prompt user to plug old adapter back in when restore fails
+        assert result3["type"] is FlowResultType.FORM
+        assert result3["step_id"] == "plug_in_new_radio"
+        assert result3["description_placeholders"] == {"device_path": "/dev/ttyUSB1234"}
+
+        # Submit retry attempt with plugged in adapter
+        result4 = await hass.config_entries.flow.async_configure(
+            result3["flow_id"],
+            user_input={},
+        )
+
+        # This adapter requires user confirmation for restore
+        assert result4["type"] is FlowResultType.FORM
+        assert result4["step_id"] == "maybe_confirm_ezsp_restore"
+
+        # Confirm destructive rewrite, but adapter is unplugged again
+        result4 = await hass.config_entries.flow.async_configure(
+            result3["flow_id"],
+            user_input={config_flow.OVERWRITE_COORDINATOR_IEEE: True},
+        )
+
+        # Prompt user to plug old adapter back in again
+        assert result4["type"] is FlowResultType.FORM
+        assert result4["step_id"] == "plug_in_new_radio"
+        assert result4["description_placeholders"] == {"device_path": "/dev/ttyUSB1234"}
+
+        # User confirms they plugged in the adapter
+        result5 = await hass.config_entries.flow.async_configure(
+            result4["flow_id"],
+            user_input={},
+        )
+
+    # Entry created successfully
+    assert result5["type"] is FlowResultType.CREATE_ENTRY
+    assert result5["data"][CONF_RADIO_TYPE] == "ezsp"
+
+    # Verify restore was attempted four times:
+    # first fail + retry for destructive dialog + failed destructive + successful retry
+    assert mock_restore_backup.call_count == 4
+
+
+@patch(f"zigpy_znp.{PROBE_FUNCTION_PATH}", AsyncMock(return_value=True))
+@patch("homeassistant.components.zha.async_setup_entry", AsyncMock(return_value=True))
+async def test_plug_in_old_radio_retry(hass: HomeAssistant, backup, mock_app) -> None:
+    """Test plug_in_old_radio step when reset fails due to unplugged adapter."""
+    entry = MockConfigEntry(
+        version=config_flow.ZhaConfigFlowHandler.VERSION,
+        domain=DOMAIN,
+        data={
+            CONF_DEVICE: {
+                CONF_DEVICE_PATH: "/dev/ttyUSB0",
+                CONF_BAUDRATE: 115200,
+                CONF_FLOW_CONTROL: None,
+            },
+            CONF_RADIO_TYPE: "ezsp",
+        },
+    )
+    entry.add_to_hass(hass)
+
+    discovery_info = UsbServiceInfo(
+        device="/dev/ttyZIGBEE",
+        pid="AAAA",
+        vid="AAAA",
+        serial_number="1234",
+        description="zigbee radio",
+        manufacturer="test",
+    )
+
+    mock_temp_radio_mgr = AsyncMock()
+    mock_temp_radio_mgr.async_reset_adapter = AsyncMock(
+        side_effect=HomeAssistantError(
+            "Failed to connect to Zigbee adapter: [Errno 2] No such file or directory"
+        )
+    )
+
+    with (
+        patch(
+            "homeassistant.components.zha.radio_manager.ZhaRadioManager._async_read_backups_from_database",
+            return_value=[backup],
+        ),
+        patch(
+            "homeassistant.components.zha.config_flow.ZhaRadioManager",
+            side_effect=[ZhaRadioManager(), mock_temp_radio_mgr, mock_temp_radio_mgr],
+        ),
+    ):
+        result_init = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": SOURCE_USB}, data=discovery_info
+        )
+
+        result_confirm = await hass.config_entries.flow.async_configure(
+            result_init["flow_id"], user_input={}
+        )
+
+        assert result_confirm["step_id"] == "choose_migration_strategy"
+
+        result_recommended = await hass.config_entries.flow.async_configure(
+            result_confirm["flow_id"],
+            user_input={"next_step_id": config_flow.MIGRATION_STRATEGY_RECOMMENDED},
+        )
+
+        # Prompt user to plug old adapter back in when reset fails
+        assert result_recommended["type"] is FlowResultType.MENU
+        assert result_recommended["step_id"] == "plug_in_old_radio"
+        assert (
+            result_recommended["description_placeholders"]["device_path"]
+            == "/dev/ttyUSB0"
+        )
+        assert result_recommended["menu_options"] == [
+            "retry_old_radio",
+            "skip_reset_old_radio",
+        ]
+
+        # Retry with unplugged adapter
+        result_retry = await hass.config_entries.flow.async_configure(
+            result_recommended["flow_id"],
+            user_input={"next_step_id": "retry_old_radio"},
+        )
+
+        # Prompt user again to plug old adapter back in
+        assert result_retry["type"] is FlowResultType.MENU
+        assert result_retry["step_id"] == "plug_in_old_radio"
+
+        # Skip resetting the old adapter
+        result_skip = await hass.config_entries.flow.async_configure(
+            result_retry["flow_id"],
+            user_input={"next_step_id": "skip_reset_old_radio"},
+        )
+
+    # Entry created successfully after skipping reset
+    assert result_skip["type"] is FlowResultType.ABORT
+    assert result_skip["reason"] == "reconfigure_successful"
+
+    # Verify reset was attempted twice: initial + retry
+    assert mock_temp_radio_mgr.async_reset_adapter.call_count == 2