Add WS command integration/wait (#142040)

* Add WS command integration/wait

* Add test

* Update homeassistant/components/websocket_api/commands.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Use helper setup.async_wait_component

* Add onboarding view

* Revert "Add onboarding view"

This reverts commit df3a1a05807ae18cac6455cf04ca0cd6bea31857.

---------

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/websocket_api/commands.py

@@ -59,7 +59,11 @@ from homeassistant.loader import (
     async_get_integration_descriptions,
     async_get_integrations,
 )
-from homeassistant.setup import async_get_loaded_integrations, async_get_setup_timings
+from homeassistant.setup import (
+    async_get_loaded_integrations,
+    async_get_setup_timings,
+    async_wait_component,
+)
 from homeassistant.util.json import format_unserializable_data
 
 from . import const, decorators, messages
@@ -98,6 +102,7 @@ def async_register_commands(
     async_reg(hass, handle_subscribe_entities)
     async_reg(hass, handle_supported_features)
     async_reg(hass, handle_integration_descriptions)
+    async_reg(hass, handle_integration_wait)
 
 
 def pong_message(iden: int) -> dict[str, Any]:
@@ -923,3 +928,21 @@ async def handle_integration_descriptions(
 ) -> None:
     """Get metadata for all brands and integrations."""
     connection.send_result(msg["id"], await async_get_integration_descriptions(hass))
+
+
+@decorators.websocket_command(
+    {
+        vol.Required("type"): "integration/wait",
+        vol.Required("domain"): str,
+    }
+)
+@decorators.async_response
+async def handle_integration_wait(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Handle wait for integration command."""
+
+    domain: str = msg["domain"]
+    connection.send_result(
+        msg["id"], {"integration_loaded": await async_wait_component(hass, domain)}
+    )

tests/components/websocket_api/test_commands.py

@@ -26,15 +26,17 @@ from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.loader import async_get_integration
-from homeassistant.setup import async_setup_component
+from homeassistant.setup import async_set_domains_to_be_loaded, async_setup_component
 from homeassistant.util.json import json_loads
 
 from tests.common import (
     MockConfigEntry,
     MockEntity,
     MockEntityPlatform,
+    MockModule,
     MockUser,
     async_mock_service,
+    mock_integration,
     mock_platform,
 )
 from tests.typing import (
@@ -2824,3 +2826,83 @@ async def test_subscribe_entities_chained_state_change(
 
     await websocket_client.close()
     await hass.async_block_till_done()
+
+
+@pytest.mark.parametrize(
+    ("domain", "result"),
+    [
+        ("config", {"integration_loaded": True}),
+        ("non_existing_domain", {"integration_loaded": False}),
+    ],
+)
+async def test_wait_integration(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    domain: str,
+    result: dict[str, Any],
+) -> None:
+    """Test we can get wait for an integration to load."""
+    assert await async_setup_component(hass, "config", {})
+    ws_client = await hass_ws_client(hass)
+
+    await ws_client.send_json_auto_id({"type": "integration/wait", "domain": domain})
+    response = await ws_client.receive_json()
+    assert response == {
+        "id": ANY,
+        "result": result,
+        "success": True,
+        "type": "result",
+    }
+
+
+async def test_wait_integration_startup(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+) -> None:
+    """Test we can get wait for an integration to load during startup."""
+    ws_client = await hass_ws_client(hass)
+
+    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
+    await ws_client.send_json_auto_id({"type": "integration/wait", "domain": "test"})
+    response = await ws_client.receive_json()
+    assert response == {
+        "id": ANY,
+        "result": {"integration_loaded": False},
+        "success": True,
+        "type": "result",
+    }
+
+    # 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
+    await ws_client.send_json_auto_id({"type": "integration/wait", "domain": "test"})
+    response = await ws_client.receive_json()
+    assert response == {
+        "id": ANY,
+        "result": {"integration_loaded": True},
+        "success": True,
+        "type": "result",
+    }
+
+    # The component has been loaded
+    assert "test" in hass.config.components