From 7fd85f54e4d46464405cae379ab3da0295d0ff49 Mon Sep 17 00:00:00 2001 From: G Johansson Date: Wed, 16 Sep 2026 10:43:05 +0200 Subject: [PATCH] Add async_retry_migration to Config Entries (#182054) Co-authored-by: Markus Tuominen <3738613+Markus98@users.noreply.github.com> --- homeassistant/config_entries.py | 24 +++++ tests/test_config_entries.py | 164 ++++++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+) diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index 94645ea1f169..80c4792156a9 100644 --- a/homeassistant/config_entries.py +++ b/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) diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py index 32adc595682c..afe1735c4f14 100644 --- a/tests/test_config_entries.py +++ b/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", [