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

@@ -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):

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")

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

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",

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(

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)
 

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,