Add onboarding view /api/onboarding/integration/wait (#142688)

homeassistant/components/onboarding/views.py

@@ -31,7 +31,12 @@ from homeassistant.helpers import area_registry as ar
 from homeassistant.helpers.backup import async_get_manager as async_get_backup_manager
 from homeassistant.helpers.system_info import async_get_system_info
 from homeassistant.helpers.translation import async_get_translations
-from homeassistant.setup import SetupPhases, async_pause_setup, async_setup_component
+from homeassistant.setup import (
+    SetupPhases,
+    async_pause_setup,
+    async_setup_component,
+    async_wait_component,
+)
 
 if TYPE_CHECKING:
     from . import OnboardingData, OnboardingStorage, OnboardingStoreData
@@ -60,6 +65,7 @@ async def async_setup(
     hass.http.register_view(BackupInfoView(data))
     hass.http.register_view(RestoreBackupView(data))
     hass.http.register_view(UploadBackupView(data))
+    hass.http.register_view(WaitIntegrationOnboardingView(data))
     await setup_cloud_views(hass, data)
 
 
@@ -298,6 +304,30 @@ class IntegrationOnboardingView(_BaseOnboardingStepView):
             return self.json({"auth_code": auth_code})
 
 
+class WaitIntegrationOnboardingView(_NoAuthBaseOnboardingView):
+    """Get backup info view."""
+
+    url = "/api/onboarding/integration/wait"
+    name = "api:onboarding:integration:wait"
+
+    @RequestDataValidator(
+        vol.Schema(
+            {
+                vol.Required("domain"): str,
+            }
+        )
+    )
+    async def post(self, request: web.Request, data: dict[str, Any]) -> web.Response:
+        """Handle wait for integration command."""
+        hass = request.app[KEY_HASS]
+        domain = data["domain"]
+        return self.json(
+            {
+                "integration_loaded": await async_wait_component(hass, domain),
+            }
+        )
+
+
 class AnalyticsOnboardingView(_BaseOnboardingStepView):
     """View to finish analytics onboarding step."""
 

tests/components/onboarding/test_views.py

@@ -21,14 +21,16 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import area_registry as ar
 from homeassistant.helpers.backup import async_initialize_backup
-from homeassistant.setup import async_setup_component
+from homeassistant.setup import async_set_domains_to_be_loaded, async_setup_component
 
 from . import mock_storage
 
 from tests.common import (
     CLIENT_ID,
     CLIENT_REDIRECT_URI,
+    MockModule,
     MockUser,
+    mock_integration,
     register_auth_provider,
 )
 from tests.test_util.aiohttp import AiohttpClientMocker
@@ -1205,3 +1207,82 @@ async def test_onboarding_cloud_status(
     assert req.status == HTTPStatus.OK
     data = await req.json()
     assert data == {"logged_in": False}
+
+
+@pytest.mark.parametrize(
+    ("domain", "expected_result"),
+    [
+        ("onboarding", {"integration_loaded": True}),
+        ("non_existing_domain", {"integration_loaded": False}),
+    ],
+)
+async def test_wait_integration(
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
+    hass_client: ClientSessionGenerator,
+    domain: str,
+    expected_result: dict[str, Any],
+) -> None:
+    """Test we can get wait for an integration to load."""
+    mock_storage(hass_storage, {"done": []})
+
+    assert await async_setup_component(hass, "onboarding", {})
+    await hass.async_block_till_done()
+
+    client = await hass_client()
+    req = await client.post("/api/onboarding/integration/wait", json={"domain": domain})
+
+    assert req.status == HTTPStatus.OK
+    data = await req.json()
+    assert data == expected_result
+
+
+async def test_wait_integration_startup(
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
+    hass_client: ClientSessionGenerator,
+) -> None:
+    """Test we can get wait for an integration to load during startup."""
+    mock_storage(hass_storage, {"done": []})
+
+    assert await async_setup_component(hass, "onboarding", {})
+    await hass.async_block_till_done()
+    client = await hass_client()
+
+    setup_stall = asyncio.Event()
+    setup_started = asyncio.Event()
+
+    async def mock_setup(hass: HomeAssistant, _) -> bool:
+        setup_started.set()
+        await setup_stall.wait()
+        return True
+
+    mock_integration(hass, MockModule("test", async_setup=mock_setup))
+
+    # The integration is not loaded, and is also not scheduled to load
+    req = await client.post("/api/onboarding/integration/wait", json={"domain": "test"})
+    assert req.status == HTTPStatus.OK
+    data = await req.json()
+    assert data == {"integration_loaded": False}
+
+    # Mark the component as scheduled to be loaded
+    async_set_domains_to_be_loaded(hass, {"test"})
+
+    # Start loading the component, including its config entries
+    hass.async_create_task(async_setup_component(hass, "test", {}))
+    await setup_started.wait()
+
+    # The component is not yet loaded
+    assert "test" not in hass.config.components
+
+    # Allow setup to proceed
+    setup_stall.set()
+
+    # The component is scheduled to load, this will block until the config entry is loaded
+    req = await client.post("/api/onboarding/integration/wait", json={"domain": "test"})
+    assert req.status == HTTPStatus.OK
+    data = await req.json()
+    assert data == {"integration_loaded": True}
+
+    # The component has been loaded
+    assert "test" in hass.config.components