From 1adfa6bbeb4e61714065d99a1169fb457d4b7d13 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 6 Aug 2023 20:25:03 -1000
Subject: [PATCH] Reduce overhead to start a config entry flow by optimizing
 fetching the handler (#97883)

---
 homeassistant/config_entries.py | 30 ++++++++++++++++++------------
 tests/test_config_entries.py    | 22 ++++++++++++++++++++++
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py
index 15fcb9a50de..1e3af81395a 100644
--- a/homeassistant/config_entries.py
+++ b/homeassistant/config_entries.py
@@ -964,10 +964,9 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager):
 
         Handler key is the domain of the component that we want to set up.
         """
-        await _load_integration(self.hass, handler_key, self._hass_config)
-        if (handler := HANDLERS.get(handler_key)) is None:
-            raise data_entry_flow.UnknownHandler
-
+        handler = await _async_get_flow_handler(
+            self.hass, handler_key, self._hass_config
+        )
         if not context or "source" not in context:
             raise KeyError("Context not set or doesn't have a source set")
 
@@ -1830,12 +1829,8 @@ class OptionsFlowManager(data_entry_flow.FlowManager):
         if entry is None:
             raise UnknownEntry(handler_key)
 
-        await _load_integration(self.hass, entry.domain, {})
-
-        if entry.domain not in HANDLERS:
-            raise data_entry_flow.UnknownHandler
-
-        return HANDLERS[entry.domain].async_get_options_flow(entry)
+        handler = await _async_get_flow_handler(self.hass, entry.domain, {})
+        return handler.async_get_options_flow(entry)
 
     async def async_finish_flow(
         self, flow: data_entry_flow.FlowHandler, result: data_entry_flow.FlowResult
@@ -2021,9 +2016,15 @@ async def support_remove_from_device(hass: HomeAssistant, domain: str) -> bool:
     return hasattr(component, "async_remove_config_entry_device")
 
 
-async def _load_integration(
+async def _async_get_flow_handler(
     hass: HomeAssistant, domain: str, hass_config: ConfigType
-) -> None:
+) -> type[ConfigFlow]:
+    """Get a flow handler for specified domain."""
+
+    # First check if there is a handler registered for the domain
+    if domain in hass.config.components and (handler := HANDLERS.get(domain)):
+        return handler
+
     try:
         integration = await loader.async_get_integration(hass, domain)
     except loader.IntegrationNotFound as err:
@@ -2042,3 +2043,8 @@ async def _load_integration(
             err,
         )
         raise data_entry_flow.UnknownHandler
+
+    if handler := HANDLERS.get(domain):
+        return handler
+
+    raise data_entry_flow.UnknownHandler
diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py
index 68e6fc59987..3485162cbb3 100644
--- a/tests/test_config_entries.py
+++ b/tests/test_config_entries.py
@@ -1544,6 +1544,28 @@ async def test_init_custom_integration(hass: HomeAssistant) -> None:
         await hass.config_entries.flow.async_init("bla", context={"source": "user"})
 
 
+async def test_init_custom_integration_with_missing_handler(
+    hass: HomeAssistant,
+) -> None:
+    """Test initializing flow for custom integration with a missing handler."""
+    integration = loader.Integration(
+        hass,
+        "custom_components.hue",
+        None,
+        {"name": "Hue", "dependencies": [], "requirements": [], "domain": "hue"},
+    )
+    mock_integration(
+        hass,
+        MockModule("hue"),
+    )
+    mock_entity_platform(hass, "config_flow.hue", None)
+    with pytest.raises(data_entry_flow.UnknownHandler), patch(
+        "homeassistant.loader.async_get_integration",
+        return_value=integration,
+    ):
+        await hass.config_entries.flow.async_init("bla", context={"source": "user"})
+
+
 async def test_support_entry_unload(hass: HomeAssistant) -> None:
     """Test unloading entry."""
     assert await config_entries.support_entry_unload(hass, "light")