Add the ability to process integration platforms on demand (#70174)

homeassistant/helpers/integration_platform.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
 import logging
 from typing import Any
 
@@ -12,6 +13,73 @@ from homeassistant.loader import async_get_integration, bind_hass
 from homeassistant.setup import ATTR_COMPONENT
 
 _LOGGER = logging.getLogger(__name__)
+DATA_INTEGRATION_PLATFORMS = "integration_platforms"
+
+
+@dataclass(frozen=True)
+class IntegrationPlatform:
+    """An integration platform."""
+
+    platform_name: str
+    process_platform: Callable[[HomeAssistant, str, Any], Awaitable[None]]
+    seen_components: set[str]
+
+
+async def _async_process_single_integration_platform(
+    hass: HomeAssistant, component_name: str, integration_platform: IntegrationPlatform
+) -> None:
+    """Process a single integration platform."""
+    if component_name in integration_platform.seen_components:
+        return
+    integration_platform.seen_components.add(component_name)
+
+    integration = await async_get_integration(hass, component_name)
+    platform_name = integration_platform.platform_name
+
+    try:
+        platform = integration.get_platform(platform_name)
+    except ImportError as err:
+        if f"{component_name}.{platform_name}" not in str(err):
+            _LOGGER.exception(
+                "Unexpected error importing %s/%s.py",
+                component_name,
+                platform_name,
+            )
+        return
+
+    try:
+        await integration_platform.process_platform(hass, component_name, platform)  # type: ignore[misc,operator] # https://github.com/python/mypy/issues/5485
+    except Exception:  # pylint: disable=broad-except
+        _LOGGER.exception(
+            "Error processing platform %s.%s", component_name, platform_name
+        )
+
+
+async def async_process_integration_platform(
+    hass: HomeAssistant, component_name: str
+) -> None:
+    """Process integration platforms on demand.
+
+    This function will load the integration platforms
+    for an integration instead of waiting for the EVENT_COMPONENT_LOADED
+    event to be fired for the integration.
+
+    When the integration will create entities before
+    it has finished setting up; call this function to ensure
+    that the integration platforms are loaded before the entities
+    are created.
+    """
+    integration_platforms: list[IntegrationPlatform] = hass.data[
+        DATA_INTEGRATION_PLATFORMS
+    ]
+    await asyncio.gather(
+        *[
+            _async_process_single_integration_platform(
+                hass, component_name, integration_platform
+            )
+            for integration_platform in integration_platforms
+        ]
+    )
 
 
 @bind_hass
@@ -22,39 +90,30 @@ async def async_process_integration_platforms(
     process_platform: Callable[[HomeAssistant, str, Any], Awaitable[None]],
 ) -> None:
     """Process a specific platform for all current and future loaded integrations."""
+    if DATA_INTEGRATION_PLATFORMS not in hass.data:
+        hass.data[DATA_INTEGRATION_PLATFORMS] = []
 
-    async def _process(component_name: str) -> None:
-        """Process component being loaded."""
-        if "." in component_name:
-            return
+        async def _async_component_loaded(event: Event) -> None:
+            """Handle a new component loaded."""
+            comp = event.data[ATTR_COMPONENT]
+            if "." not in comp:
+                await async_process_integration_platform(hass, comp)
 
-        integration = await async_get_integration(hass, component_name)
+        hass.bus.async_listen(EVENT_COMPONENT_LOADED, _async_component_loaded)
 
-        try:
-            platform = integration.get_platform(platform_name)
-        except ImportError as err:
-            if f"{component_name}.{platform_name}" not in str(err):
-                _LOGGER.exception(
-                    "Unexpected error importing %s/%s.py",
-                    component_name,
-                    platform_name,
+    integration_platforms: list[IntegrationPlatform] = hass.data[
+        DATA_INTEGRATION_PLATFORMS
+    ]
+    integration_platform = IntegrationPlatform(platform_name, process_platform, set())
+    integration_platforms.append(integration_platform)
+    if top_level_components := (
+        comp for comp in hass.config.components if "." not in comp
+    ):
+        await asyncio.gather(
+            *[
+                _async_process_single_integration_platform(
+                    hass, comp, integration_platform
                 )
-            return
-
-        try:
-            await process_platform(hass, component_name, platform)
-        except Exception:  # pylint: disable=broad-except
-            _LOGGER.exception(
-                "Error processing platform %s.%s", component_name, platform_name
-            )
-
-    async def async_component_loaded(event: Event) -> None:
-        """Handle a new component loaded."""
-        await _process(event.data[ATTR_COMPONENT])
-
-    hass.bus.async_listen(EVENT_COMPONENT_LOADED, async_component_loaded)
-
-    tasks = [_process(comp) for comp in hass.config.components]
-
-    if tasks:
-        await asyncio.gather(*tasks)
+                for comp in top_level_components
+            ]
+        )

tests/components/intent/test_init.py

@@ -60,6 +60,7 @@ async def test_cover_intents_loading(hass):
         )
 
     assert await async_setup_component(hass, "cover", {})
+    await hass.async_block_till_done()
 
     hass.states.async_set("cover.garage_door", "closed")
     calls = async_mock_service(hass, "cover", SERVICE_OPEN_COVER)
@@ -81,6 +82,7 @@ async def test_turn_on_intent(hass):
     """Test HassTurnOn intent."""
     result = await async_setup_component(hass, "homeassistant", {})
     result = await async_setup_component(hass, "intent", {})
+    await hass.async_block_till_done()
     assert result
 
     hass.states.async_set("light.test_light", "off")

tests/components/lovelace/test_system_health.py

@@ -24,6 +24,7 @@ async def test_system_health_info_storage(hass, hass_storage):
         "data": {"config": {"resources": [], "views": []}},
     }
     assert await async_setup_component(hass, "lovelace", {})
+    await hass.async_block_till_done()
     info = await get_system_health_info(hass, "lovelace")
     assert info == {"dashboards": 1, "mode": "storage", "resources": 0, "views": 0}
 
@@ -32,6 +33,7 @@ async def test_system_health_info_yaml(hass):
     """Test system health info endpoint."""
     assert await async_setup_component(hass, "system_health", {})
     assert await async_setup_component(hass, "lovelace", {"lovelace": {"mode": "YAML"}})
+    await hass.async_block_till_done()
     with patch(
         "homeassistant.components.lovelace.dashboard.load_yaml",
         return_value={"views": [{"cards": []}]},
@@ -44,6 +46,7 @@ async def test_system_health_info_yaml_not_found(hass):
     """Test system health info endpoint."""
     assert await async_setup_component(hass, "system_health", {})
     assert await async_setup_component(hass, "lovelace", {"lovelace": {"mode": "YAML"}})
+    await hass.async_block_till_done()
     info = await get_system_health_info(hass, "lovelace")
     assert info == {
         "dashboards": 1,

tests/helpers/test_integration_platform.py

@@ -35,3 +35,9 @@ async def test_process_integration_platforms(hass):
     assert len(processed) == 2
     assert processed[1][0] == "event"
     assert processed[1][1] == event_platform
+
+    # Verify we only process the platform once if we call it manually
+    await hass.helpers.integration_platform.async_process_integration_platform(
+        hass, "event"
+    )
+    assert len(processed) == 2