From bc08d657d4bdde8ef0a74cdfea042b8fe07f490e Mon Sep 17 00:00:00 2001
From: Erik <erik@montnemery.com>
Date: Wed, 2 Apr 2025 08:34:21 +0200
Subject: [PATCH] Add WS command integration/wait

---
 .../components/websocket_api/commands.py      | 31 ++++++++++++++++++-
 .../components/websocket_api/test_commands.py | 27 ++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/websocket_api/commands.py b/homeassistant/components/websocket_api/commands.py
index 4a360b4a43c..35bd6b1cd7c 100644
--- a/homeassistant/components/websocket_api/commands.py
+++ b/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 (
+    DATA_SETUP_DONE,
+    async_get_loaded_integrations,
+    async_get_setup_timings,
+)
 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,27 @@ 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 = msg["domain"]
+    setup_done = hass.data.get(DATA_SETUP_DONE, {})
+
+    # Wait for the integration to be set up
+    if setup_future := setup_done.get(domain):
+        await setup_future
+
+    connection.send_result(
+        msg["id"], {"integration_loaded": domain in hass.config.components}
+    )
diff --git a/tests/components/websocket_api/test_commands.py b/tests/components/websocket_api/test_commands.py
index f03673048c0..a4f47fce594 100644
--- a/tests/components/websocket_api/test_commands.py
+++ b/tests/components/websocket_api/test_commands.py
@@ -2824,3 +2824,30 @@ 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",
+    }