From 9940f51b95b7ba3423cd6079bd83ecf325bca7c2 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 13 Mar 2024 18:13:40 -1000 Subject: [PATCH] Avoid pre-importing config_flows if the integration does not support migration (#113369) * Avoid pre-importing config_flows if the integration does support migration Currently we pre-import the config flow module if it exists since setting up the config entry required comparing the versions found in the config_flow.py. We can avoid the pre-import if the integration does not support async_migrate_entry which means we avoid loading many config flows in memory at startup. * cover * fix missing block * do not call directly * its too fast now, the test gets more along * Update homeassistant/loader.py --- homeassistant/config_entries.py | 22 +++++-- homeassistant/loader.py | 23 ++++++-- tests/components/jellyfin/test_init.py | 1 + tests/components/sonarr/test_init.py | 4 +- .../unifiprotect/test_config_flow.py | 57 +++++++++++++------ tests/test_config_entries.py | 3 + tests/test_loader.py | 21 +++---- 7 files changed, 89 insertions(+), 42 deletions(-) diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index 4b2eb81c870..b7623e9c8d7 100644 --- a/homeassistant/config_entries.py +++ b/homeassistant/config_entries.py @@ -349,6 +349,9 @@ class ConfigEntry: # Supports remove device self.supports_remove_device: bool | None = None + # Supports migrate + self.supports_migrate: bool | None = None + # Supports options self._supports_options: bool | None = None @@ -491,6 +494,7 @@ class ConfigEntry: self.supports_remove_device = await support_remove_from_device( hass, self.domain ) + try: component = await integration.async_get_component() except ImportError as err: @@ -506,7 +510,12 @@ class ConfigEntry: ) return - if domain_is_integration: + if self.supports_migrate is None: + self.supports_migrate = hasattr(component, "async_migrate_entry") + + if domain_is_integration and self.supports_migrate: + # Avoid loading the config_flow module unless we need to check + # the version to see if we need to migrate try: await integration.async_get_platforms(("config_flow",)) except ImportError as err: @@ -787,11 +796,7 @@ class ConfigEntry: if same_major_version and self.minor_version == handler.MINOR_VERSION: return True - if not (integration := self._integration_for_domain): - integration = await loader.async_get_integration(hass, self.domain) - component = await integration.async_get_component() - supports_migrate = hasattr(component, "async_migrate_entry") - if not supports_migrate: + if not self.supports_migrate: if same_major_version: return True _LOGGER.error( @@ -801,6 +806,11 @@ class ConfigEntry: ) return False + if not (integration := self._integration_for_domain): + integration = await loader.async_get_integration(hass, self.domain) + + component = await integration.async_get_component() + try: result = await component.async_migrate_entry(hass, self) if not isinstance(result, bool): diff --git a/homeassistant/loader.py b/homeassistant/loader.py index 8e44ca38d77..9c961fd7f67 100644 --- a/homeassistant/loader.py +++ b/homeassistant/loader.py @@ -961,10 +961,7 @@ class Integration: # Some integrations fail on import because they call functions incorrectly. # So we do it before validating config to catch these errors. - load_executor = self.import_executor and ( - self.pkg_path not in sys.modules - or (self.config_flow and f"{self.pkg_path}.config_flow" not in sys.modules) - ) + load_executor = self.import_executor and self.pkg_path not in sys.modules if not load_executor: comp = self._get_component() if debug: @@ -1051,6 +1048,12 @@ class Integration: if preload_platforms: for platform_name in self.platforms_exists(self._platforms_to_preload): + if ( + platform_name == "config_flow" + and not async_config_flow_needs_preload(cache[domain]) + ): + continue + with suppress(ImportError): self.get_platform(platform_name) @@ -1684,3 +1687,15 @@ def async_suggest_report_issue( ) return f"create a bug report at {issue_tracker}" + + +@callback +def async_config_flow_needs_preload(component: ComponentProtocol) -> bool: + """Test if a config_flow for a component needs to be preloaded. + + Currently we need to preload a the config flow if the integration + has a config flow and the component has an async_migrate_entry method + because it means that config_entries will always have to load + it to check if it needs to be migrated. + """ + return hasattr(component, "async_migrate_entry") diff --git a/tests/components/jellyfin/test_init.py b/tests/components/jellyfin/test_init.py index 9bfe37f9874..6e6a0f7219b 100644 --- a/tests/components/jellyfin/test_init.py +++ b/tests/components/jellyfin/test_init.py @@ -67,6 +67,7 @@ async def test_invalid_auth( mock_config_entry.add_to_hass(hass) assert not await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() flows = hass.config_entries.flow.async_progress() assert len(flows) == 1 diff --git a/tests/components/sonarr/test_init.py b/tests/components/sonarr/test_init.py index 9f512c11074..e663139d33c 100644 --- a/tests/components/sonarr/test_init.py +++ b/tests/components/sonarr/test_init.py @@ -105,7 +105,9 @@ async def test_migrate_config_entry(hass: HomeAssistant) -> None: assert entry.version == 1 assert not entry.unique_id - await entry.async_migrate(hass) + with patch("homeassistant.components.sonarr.async_setup_entry", return_value=True): + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() assert entry.data == { CONF_API_KEY: "MOCK_API_KEY", diff --git a/tests/components/unifiprotect/test_config_flow.py b/tests/components/unifiprotect/test_config_flow.py index 6dfbb907097..a5fbdb8493f 100644 --- a/tests/components/unifiprotect/test_config_flow.py +++ b/tests/components/unifiprotect/test_config_flow.py @@ -219,25 +219,31 @@ async def test_form_reauth_auth( ) mock_config.add_to_hass(hass) - result = await hass.config_entries.flow.async_init( - DOMAIN, - context={ - "source": config_entries.SOURCE_REAUTH, - "entry_id": mock_config.entry_id, - }, - ) - assert result["type"] == FlowResultType.FORM - assert not result["errors"] - flows = hass.config_entries.flow.async_progress_by_handler(DOMAIN) - assert flows[0]["context"]["title_placeholders"] == { - "ip_address": "1.1.1.1", - "name": "Mock Title", - } - with patch( "homeassistant.components.unifiprotect.config_flow.ProtectApiClient.get_bootstrap", side_effect=NotAuthorized, + ), patch( + "homeassistant.components.unifiprotect.async_setup", + return_value=True, + ) as mock_setup, patch( + "homeassistant.components.unifiprotect.async_setup_entry", + return_value=True, ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={ + "source": config_entries.SOURCE_REAUTH, + "entry_id": mock_config.entry_id, + }, + ) + assert result["type"] == FlowResultType.FORM + assert not result["errors"] + flows = hass.config_entries.flow.async_progress_by_handler(DOMAIN) + assert flows[0]["context"]["title_placeholders"] == { + "ip_address": "1.1.1.1", + "name": "Mock Title", + } + result2 = await hass.config_entries.flow.async_configure( result["flow_id"], { @@ -257,7 +263,10 @@ async def test_form_reauth_auth( ), patch( "homeassistant.components.unifiprotect.async_setup", return_value=True, - ) as mock_setup: + ) as mock_setup, patch( + "homeassistant.components.unifiprotect.async_setup_entry", + return_value=True, + ): result3 = await hass.config_entries.flow.async_configure( result2["flow_id"], { @@ -768,6 +777,20 @@ async def test_discovered_by_unifi_discovery_direct_connect_on_different_interfa ) mock_config.add_to_hass(hass) + with patch( + "homeassistant.components.unifiprotect.async_setup_entry", + return_value=True, + ) as mock_setup_entry, patch( + "homeassistant.components.unifiprotect.async_setup", + return_value=True, + ) as mock_setup: + await hass.config_entries.async_setup(mock_config.entry_id) + await hass.async_block_till_done() + + assert mock_config.state == config_entries.ConfigEntryState.LOADED + assert len(mock_setup_entry.mock_calls) == 1 + assert len(mock_setup.mock_calls) == 1 + other_ip_dict = UNIFI_DISCOVERY_DICT.copy() other_ip_dict["source_ip"] = "127.0.0.2" other_ip_dict["direct_connect_domain"] = "nomatchsameip.ui.direct" @@ -823,7 +846,7 @@ async def test_discovered_by_unifi_discovery_direct_connect_on_different_interfa "verify_ssl": True, } assert len(mock_setup_entry.mock_calls) == 1 - assert len(mock_setup.mock_calls) == 1 + assert len(mock_setup.mock_calls) == 0 async def test_discovered_by_unifi_discovery_direct_connect_on_different_interface_resolver_no_result( diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py index 4ada764a9ba..18022ae3a21 100644 --- a/tests/test_config_entries.py +++ b/tests/test_config_entries.py @@ -117,6 +117,7 @@ async def test_call_setup_entry(hass: HomeAssistant) -> None: assert len(mock_setup_entry.mock_calls) == 1 assert entry.state is config_entries.ConfigEntryState.LOADED assert entry.supports_unload + assert entry.supports_migrate async def test_call_setup_entry_without_reload_support(hass: HomeAssistant) -> None: @@ -146,6 +147,7 @@ async def test_call_setup_entry_without_reload_support(hass: HomeAssistant) -> N assert len(mock_setup_entry.mock_calls) == 1 assert entry.state is config_entries.ConfigEntryState.LOADED assert not entry.supports_unload + assert entry.supports_migrate @pytest.mark.parametrize(("major_version", "minor_version"), [(2, 1), (1, 2), (2, 2)]) @@ -289,6 +291,7 @@ async def test_call_async_migrate_entry_failure_not_supported( ) entry.add_to_hass(hass) assert not entry.supports_unload + entry.supports_migrate = True mock_setup_entry = AsyncMock(return_value=True) diff --git a/tests/test_loader.py b/tests/test_loader.py index a2868976876..5bdf93758b0 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -1204,31 +1204,21 @@ async def test_async_get_component_loads_loop_if_already_in_sys_modules( assert "test_package_loaded_executor" not in hass.config.components assert "test_package_loaded_executor.config_flow" not in hass.config.components - config_flow_module_name = f"{integration.pkg_path}.config_flow" - module_mock = MagicMock(__file__="__init__.py") - config_flow_module_mock = MagicMock(__file__="config_flow.py") + module_mock = object() def import_module(name: str) -> Any: if name == integration.pkg_path: return module_mock - if name == config_flow_module_name: - return config_flow_module_mock raise ImportError - modules_without_config_flow = { - k: v for k, v in sys.modules.items() if k != config_flow_module_name - } with patch.dict( "sys.modules", - {**modules_without_config_flow, integration.pkg_path: module_mock}, + {**sys.modules, integration.pkg_path: module_mock}, clear=True, ), patch("homeassistant.loader.importlib.import_module", import_module): module = await integration.async_get_component() - # The config flow is missing so we should load - # in the executor - assert "loaded_executor=True" in caplog.text - assert "loaded_executor=False" not in caplog.text + assert "loaded_executor=False" in caplog.text assert module is module_mock caplog.clear() @@ -1236,7 +1226,6 @@ async def test_async_get_component_loads_loop_if_already_in_sys_modules( "sys.modules", { integration.pkg_path: module_mock, - config_flow_module_name: config_flow_module_mock, }, ), patch("homeassistant.loader.importlib.import_module", import_module): module = await integration.async_get_component() @@ -1246,6 +1235,10 @@ async def test_async_get_component_loads_loop_if_already_in_sys_modules( assert "loaded_executor" not in caplog.text assert module is module_mock + # The integration does not implement async_migrate_entry so it + # should not be preloaded + assert integration.get_platform_cached("config_flow") is None + async def test_async_get_component_concurrent_loads( hass: HomeAssistant,