Ensure dependencies are awaited correctly when setting up integrations (#91454)

* Do not wait

* Correct tests

* Manage after dependencies stage 1

* test bootstrap dependencies

* Assert log the dependenciy is waited for

* Improve docstrings

* Assert outside callback

* Patch async_get_integrations

* Revert changes made to snips integration

* Undo changes to mqtt_statestream

homeassistant/bootstrap.py

@@ -629,6 +629,9 @@ async def _async_set_up_integrations(
         - stage_1_domains
     )
 
+    # Enables after dependencies when setting up stage 1 domains
+    async_set_domains_to_be_loaded(hass, stage_1_domains)
+
     # Start setup
     if stage_1_domains:
         _LOGGER.info("Setting up stage 1: %s", stage_1_domains)
@@ -640,7 +643,7 @@ async def _async_set_up_integrations(
         except asyncio.TimeoutError:
             _LOGGER.warning("Setup timed out for stage 1 - moving forward")
 
-    # Enables after dependencies
+    # Add after dependencies when setting up stage 2 domains
     async_set_domains_to_be_loaded(hass, stage_2_domains)
 
     if stage_2_domains:

homeassistant/setup.py

@@ -67,7 +67,8 @@ def async_set_domains_to_be_loaded(hass: core.HomeAssistant, domains: set[str])
      - Properly handle after_dependencies.
      - Keep track of domains which will load but have not yet finished loading
     """
-    hass.data[DATA_SETUP_DONE] = {domain: asyncio.Event() for domain in domains}
+    hass.data.setdefault(DATA_SETUP_DONE, {})
+    hass.data[DATA_SETUP_DONE].update({domain: asyncio.Event() for domain in domains})
 
 
 def setup_component(hass: core.HomeAssistant, domain: str, config: ConfigType) -> bool:

tests/test_bootstrap.py

@@ -1,6 +1,6 @@
 """Test the bootstrapping."""
 import asyncio
-from collections.abc import Generator
+from collections.abc import Generator, Iterable
 import glob
 import os
 from typing import Any
@@ -10,12 +10,16 @@ import pytest
 
 from homeassistant import bootstrap, runner
 import homeassistant.config as config_util
+from homeassistant.config_entries import HANDLERS, ConfigEntry
 from homeassistant.const import SIGNAL_BOOTSTRAP_INTEGRATIONS
 from homeassistant.core import HomeAssistant, async_get_hass, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.typing import ConfigType
+from homeassistant.loader import Integration
 
 from .common import (
+    MockConfigEntry,
     MockModule,
     MockPlatform,
     get_test_config_dir,
@@ -825,3 +829,131 @@ async def test_bootstrap_empty_integrations(
     """Test setting up an empty integrations does not raise."""
     await bootstrap.async_setup_multi_components(hass, set(), {})
     await hass.async_block_till_done()
+
+
+@pytest.mark.parametrize("integration", ["mqtt_eventstream", "mqtt_statestream"])
+@pytest.mark.parametrize("load_registries", [False])
+async def test_bootstrap_dependencies(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+    integration: str,
+) -> None:
+    """Test dependencies are set up correctly,."""
+
+    # Prepare MQTT config entry
+    @HANDLERS.register("mqtt")
+    class MockConfigFlow:
+        """Mock the MQTT config flow."""
+
+        VERSION = 1
+
+    entry = MockConfigEntry(domain="mqtt", data={"broker": "test-broker"})
+    entry.add_to_hass(hass)
+
+    calls: list[str] = []
+    assertions: list[bool] = []
+
+    async def async_mqtt_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+        """Assert the mqtt config entry was set up."""
+        calls.append("mqtt")
+        # assert the integration is not yet set up
+        assertions.append(hass.data["setup_done"][integration].is_set() is False)
+        assertions.append(
+            all(
+                dependency in hass.config.components
+                for dependency in integrations[integration]["dependencies"]
+            )
+        )
+        assertions.append(integration not in hass.config.components)
+        return True
+
+    async def async_integration_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Assert the mqtt config entry was set up."""
+        calls.append(integration)
+        # assert mqtt was already set up
+        assertions.append(
+            "mqtt" not in hass.data["setup_done"]
+            or hass.data["setup_done"]["mqtt"].is_set()
+        )
+        assertions.append("mqtt" in hass.config.components)
+        return True
+
+    mqtt_integration = mock_integration(
+        hass,
+        MockModule(
+            "mqtt",
+            async_setup_entry=async_mqtt_setup_entry,
+            dependencies=["file_upload", "http"],
+        ),
+    )
+    mqtt_integration._import_platform = Mock()
+    # mqtt_integration.async_migrate = AsyncMock(return_value=False)
+
+    integrations = {
+        "mqtt": {
+            "dependencies": {"file_upload", "http"},
+            "integration": mqtt_integration,
+        },
+        "mqtt_eventstream": {
+            "dependencies": {"mqtt"},
+            "integration": mock_integration(
+                hass,
+                MockModule(
+                    "mqtt_eventstream",
+                    async_setup=async_integration_setup,
+                    dependencies=["mqtt"],
+                ),
+            ),
+        },
+        "mqtt_statestream": {
+            "dependencies": {"mqtt"},
+            "integration": mock_integration(
+                hass,
+                MockModule(
+                    "mqtt_statestream",
+                    async_setup=async_integration_setup,
+                    dependencies=["mqtt"],
+                ),
+            ),
+        },
+        "file_upload": {
+            "dependencies": {"http"},
+            "integration": mock_integration(
+                hass,
+                MockModule(
+                    "file_upload",
+                    dependencies=["http"],
+                ),
+            ),
+        },
+        "http": {
+            "dependencies": set(),
+            "integration": mock_integration(
+                hass,
+                MockModule("http", dependencies=[]),
+            ),
+        },
+    }
+
+    async def mock_async_get_integrations(
+        hass: HomeAssistant, domains: Iterable[str]
+    ) -> dict[str, Integration | Exception]:
+        """Mock integrations."""
+        return {domain: integrations[domain]["integration"] for domain in domains}
+
+    with patch(
+        "homeassistant.setup.loader.async_get_integrations",
+        side_effect=mock_async_get_integrations,
+    ), patch("homeassistant.config.async_process_component_config", return_value={}):
+        bootstrap.async_set_domains_to_be_loaded(hass, {integration})
+        await bootstrap.async_setup_multi_components(hass, {integration}, {})
+        await hass.async_block_till_done()
+
+    for assertion in assertions:
+        assert assertion
+
+    assert calls == ["mqtt", integration]
+
+    assert (
+        f"Dependency {integration} will wait for dependencies ['mqtt']" in caplog.text
+    )