Add async_retry_migration to Config Entries (#182054)

Co-authored-by: Markus Tuominen <3738613+Markus98@users.noreply.github.com>

homeassistant/config_entries.py

@@ -2527,6 +2527,30 @@ class ConfigEntries:
             entry.state is ConfigEntryState.LOADED  # type: ignore[comparison-overlap]
         )
 
+    async def async_retry_migration(self, entry_id: str) -> None:
+        """Retry migration for a config entry.
+
+        This is only intended for repair flows created to handle
+        non-recoverable migration errors.
+        """
+        entry = self.async_get_known_entry(entry_id)
+        if entry.state is not ConfigEntryState.MIGRATION_ERROR:
+            raise OperationNotAllowed(
+                f"The config entry '{entry.title}' ({entry.domain}) with entry_id"
+                f" '{entry.entry_id}' cannot retry the migration as it is not in the"
+                f" state {ConfigEntryState.MIGRATION_ERROR} but is in the state {entry.state}"
+            )
+        if entry.disabled_by:
+            raise OperationNotAllowed(
+                f"The config entry '{entry.title}' ({entry.domain}) with entry_id"
+                f" '{entry.entry_id}' cannot retry the migration as it is disabled by"
+                f" {entry.disabled_by}. Please enable the config entry and retry."
+            )
+
+        # Config entry was never loaded so we can set state and start setup to try again
+        entry._async_set_state(self.hass, ConfigEntryState.NOT_LOADED, None)  # noqa: SLF001
+        await self.async_setup(entry_id)
+
     async def async_unload(self, entry_id: str, _lock: bool = True) -> bool:
         """Unload a config entry."""
         entry = self.async_get_known_entry(entry_id)

tests/test_config_entries.py

@@ -3729,6 +3729,170 @@ async def test_entry_reload_not_loaded(
     assert entry.state is config_entries.ConfigEntryState.LOADED
 
 
+async def test_async_retry_migration(
+    hass: HomeAssistant,
+    manager: config_entries.ConfigEntries,
+) -> None:
+    """Test that we can recover from a migration error."""
+    entry = MockConfigEntry(domain="comp")
+    entry.add_to_hass(hass)
+
+    async_setup = AsyncMock(return_value=True)
+    async_setup_entry = AsyncMock(return_value=True)
+    async_unload_entry = AsyncMock(return_value=True)
+    async_migrate_entry = AsyncMock(side_effect=[False, False, True])
+
+    mock_integration(
+        hass,
+        MockModule(
+            "comp",
+            async_setup=async_setup,
+            async_setup_entry=async_setup_entry,
+            async_unload_entry=async_unload_entry,
+            async_migrate_entry=async_migrate_entry,
+        ),
+    )
+    mock_platform(hass, "comp.config_flow", None)
+
+    class TestFlow(config_entries.ConfigFlow):
+        """Test flow."""
+
+        VERSION = 2
+        MINOR_VERSION = 1
+
+        async def async_step_user(
+            self, user_input: dict[str, Any] | None = None
+        ) -> FlowResult:
+            """Test user step."""
+            return self.async_create_entry(title="title", data={})
+
+    with mock_config_flow("comp", TestFlow):
+        result = await async_setup_component(hass, "comp", {})
+        await hass.async_block_till_done()
+
+    assert result is True
+    assert entry.state is config_entries.ConfigEntryState.MIGRATION_ERROR
+
+    with mock_config_flow("comp", TestFlow):
+        await manager.async_retry_migration(entry.entry_id)
+
+    assert len(async_unload_entry.mock_calls) == 0
+    assert len(async_setup.mock_calls) == 1
+    assert len(async_setup_entry.mock_calls) == 0
+    assert len(async_migrate_entry.mock_calls) == 2
+    assert entry.version == 1
+    assert entry.state is config_entries.ConfigEntryState.MIGRATION_ERROR
+
+    with mock_config_flow("comp", TestFlow):
+        await manager.async_retry_migration(entry.entry_id)
+
+    assert len(async_unload_entry.mock_calls) == 0
+    assert len(async_setup.mock_calls) == 1
+    assert len(async_setup_entry.mock_calls) == 1
+    assert len(async_migrate_entry.mock_calls) == 3
+    assert entry.state is config_entries.ConfigEntryState.LOADED
+
+
+@pytest.mark.parametrize(
+    "state",
+    [
+        config_entries.ConfigEntryState.NOT_LOADED,
+        config_entries.ConfigEntryState.LOADED,
+        config_entries.ConfigEntryState.SETUP_ERROR,
+        config_entries.ConfigEntryState.FAILED_UNLOAD,
+    ],
+)
+async def test_async_retry_migration_fails(
+    hass: HomeAssistant,
+    manager: config_entries.ConfigEntries,
+    state: config_entries.ConfigEntryState,
+) -> None:
+    """Test we can't use async_retry_migration.
+
+    With other states than `MIGRATION_ERROR`.
+    """
+    entry = MockConfigEntry(domain="comp", state=state)
+    entry.add_to_hass(hass)
+
+    with pytest.raises(config_entries.OperationNotAllowed, match=str(state)):
+        await manager.async_retry_migration(entry.entry_id)
+
+    assert entry.state is state
+
+
+async def test_async_retry_migration_on_disabled_entry(
+    hass: HomeAssistant,
+    manager: config_entries.ConfigEntries,
+) -> None:
+    """Test we can't use async_retry_migration.
+
+    On a disabled entry that is in `MIGRATION_ERROR` state.
+    """
+    entry = MockConfigEntry(domain="comp")
+    entry.add_to_hass(hass)
+
+    async_setup = AsyncMock(return_value=True)
+    async_setup_entry = AsyncMock(return_value=True)
+    async_unload_entry = AsyncMock(return_value=True)
+    async_migrate_entry = AsyncMock(side_effect=[False, True])
+
+    mock_integration(
+        hass,
+        MockModule(
+            "comp",
+            async_setup=async_setup,
+            async_setup_entry=async_setup_entry,
+            async_unload_entry=async_unload_entry,
+            async_migrate_entry=async_migrate_entry,
+        ),
+    )
+    mock_platform(hass, "comp.config_flow", None)
+
+    class TestFlow(config_entries.ConfigFlow):
+        """Test flow."""
+
+        VERSION = 2
+        MINOR_VERSION = 1
+
+        async def async_step_user(
+            self, user_input: dict[str, Any] | None = None
+        ) -> FlowResult:
+            """Test user step."""
+            return self.async_create_entry(title="title", data={})
+
+    with mock_config_flow("comp", TestFlow):
+        result = await async_setup_component(hass, "comp", {})
+        await hass.async_block_till_done()
+
+    assert result is True
+    assert entry.state is config_entries.ConfigEntryState.MIGRATION_ERROR
+
+    with pytest.raises(config_entries.OperationNotAllowed):
+        await manager.async_set_disabled_by(
+            entry.entry_id, disabled_by=config_entries.ConfigEntryDisabler.USER
+        )
+    # Disabled by user is set even if it's not allowed to reload after
+    assert entry.disabled_by is config_entries.ConfigEntryDisabler.USER
+
+    with pytest.raises(
+        config_entries.OperationNotAllowed,
+        match=(
+            " cannot retry the migration as it is disabled by user."
+            " Please enable the config entry and retry."
+        ),
+    ):
+        await manager.async_retry_migration(entry.entry_id)
+
+    with pytest.raises(config_entries.OperationNotAllowed):
+        await manager.async_set_disabled_by(entry.entry_id, disabled_by=None)
+    assert entry.disabled_by is None
+
+    with mock_config_flow("comp", TestFlow):
+        await manager.async_retry_migration(entry.entry_id)
+
+    assert entry.state is config_entries.ConfigEntryState.LOADED
+
+
 @pytest.mark.parametrize(
     "state",
     [