Add subscribe preview feature endpoint to labs (#157976)

homeassistant/components/labs/__init__.py

@@ -7,11 +7,10 @@ in the Home Assistant Labs UI for users to enable or disable.
 
 from __future__ import annotations
 
-from collections.abc import Callable
 import logging
 
 from homeassistant.const import EVENT_LABS_UPDATED
-from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.generated.labs import LABS_PREVIEW_FEATURES
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.storage import Store
@@ -19,6 +18,7 @@ from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import async_get_custom_components
 
 from .const import DOMAIN, LABS_DATA, STORAGE_KEY, STORAGE_VERSION
+from .helpers import async_is_preview_feature_enabled, async_listen
 from .models import (
     EventLabsUpdatedData,
     LabPreviewFeature,
@@ -135,55 +135,3 @@ async def _async_scan_all_preview_features(
 
     _LOGGER.debug("Loaded %d total lab preview features", len(preview_features))
     return preview_features
-
-
-@callback
-def async_is_preview_feature_enabled(
-    hass: HomeAssistant, domain: str, preview_feature: str
-) -> bool:
-    """Check if a lab preview feature is enabled.
-
-    Args:
-        hass: HomeAssistant instance
-        domain: Integration domain
-        preview_feature: Preview feature name
-
-    Returns:
-        True if the preview feature is enabled, False otherwise
-    """
-    if LABS_DATA not in hass.data:
-        return False
-
-    labs_data = hass.data[LABS_DATA]
-    return (domain, preview_feature) in labs_data.data.preview_feature_status
-
-
-@callback
-def async_listen(
-    hass: HomeAssistant,
-    domain: str,
-    preview_feature: str,
-    listener: Callable[[], None],
-) -> Callable[[], None]:
-    """Listen for changes to a specific preview feature.
-
-    Args:
-        hass: HomeAssistant instance
-        domain: Integration domain
-        preview_feature: Preview feature name
-        listener: Callback to invoke when the preview feature is toggled
-
-    Returns:
-        Callable to unsubscribe from the listener
-    """
-
-    @callback
-    def _async_feature_updated(event: Event[EventLabsUpdatedData]) -> None:
-        """Handle labs feature update event."""
-        if (
-            event.data["domain"] == domain
-            and event.data["preview_feature"] == preview_feature
-        ):
-            listener()
-
-    return hass.bus.async_listen(EVENT_LABS_UPDATED, _async_feature_updated)

homeassistant/components/labs/helpers.py

@@ -0,0 +1,63 @@
+"""Helper functions for the Home Assistant Labs integration."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from homeassistant.const import EVENT_LABS_UPDATED
+from homeassistant.core import Event, HomeAssistant, callback
+
+from .const import LABS_DATA
+from .models import EventLabsUpdatedData
+
+
+@callback
+def async_is_preview_feature_enabled(
+    hass: HomeAssistant, domain: str, preview_feature: str
+) -> bool:
+    """Check if a lab preview feature is enabled.
+
+    Args:
+        hass: HomeAssistant instance
+        domain: Integration domain
+        preview_feature: Preview feature name
+
+    Returns:
+        True if the preview feature is enabled, False otherwise
+    """
+    if LABS_DATA not in hass.data:
+        return False
+
+    labs_data = hass.data[LABS_DATA]
+    return (domain, preview_feature) in labs_data.data.preview_feature_status
+
+
+@callback
+def async_listen(
+    hass: HomeAssistant,
+    domain: str,
+    preview_feature: str,
+    listener: Callable[[], None],
+) -> Callable[[], None]:
+    """Listen for changes to a specific preview feature.
+
+    Args:
+        hass: HomeAssistant instance
+        domain: Integration domain
+        preview_feature: Preview feature name
+        listener: Callback to invoke when the preview feature is toggled
+
+    Returns:
+        Callable to unsubscribe from the listener
+    """
+
+    @callback
+    def _async_feature_updated(event: Event[EventLabsUpdatedData]) -> None:
+        """Handle labs feature update event."""
+        if (
+            event.data["domain"] == domain
+            and event.data["preview_feature"] == preview_feature
+        ):
+            listener()
+
+    return hass.bus.async_listen(EVENT_LABS_UPDATED, _async_feature_updated)

homeassistant/components/labs/websocket_api.py

@@ -12,6 +12,7 @@ from homeassistant.const import EVENT_LABS_UPDATED
 from homeassistant.core import HomeAssistant, callback
 
 from .const import LABS_DATA
+from .helpers import async_is_preview_feature_enabled, async_listen
 from .models import EventLabsUpdatedData
 
 
@@ -20,6 +21,7 @@ def async_setup(hass: HomeAssistant) -> None:
     """Set up the number websocket API."""
     websocket_api.async_register_command(hass, websocket_list_preview_features)
     websocket_api.async_register_command(hass, websocket_update_preview_feature)
+    websocket_api.async_register_command(hass, websocket_subscribe_feature)
 
 
 @callback
@@ -108,3 +110,52 @@ async def websocket_update_preview_feature(
     hass.bus.async_fire(EVENT_LABS_UPDATED, event_data)
 
     connection.send_result(msg["id"])
+
+
+@callback
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "labs/subscribe",
+        vol.Required("domain"): str,
+        vol.Required("preview_feature"): str,
+    }
+)
+def websocket_subscribe_feature(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
+    """Subscribe to a specific lab preview feature updates."""
+    domain = msg["domain"]
+    preview_feature_key = msg["preview_feature"]
+    labs_data = hass.data[LABS_DATA]
+
+    preview_feature_id = f"{domain}.{preview_feature_key}"
+
+    if preview_feature_id not in labs_data.preview_features:
+        connection.send_error(
+            msg["id"],
+            websocket_api.ERR_NOT_FOUND,
+            f"Preview feature {preview_feature_id} not found",
+        )
+        return
+
+    preview_feature = labs_data.preview_features[preview_feature_id]
+
+    @callback
+    def send_event() -> None:
+        """Send feature state to client."""
+        enabled = async_is_preview_feature_enabled(hass, domain, preview_feature_key)
+        connection.send_message(
+            websocket_api.event_message(
+                msg["id"],
+                preview_feature.to_dict(enabled=enabled),
+            )
+        )
+
+    connection.subscriptions[msg["id"]] = async_listen(
+        hass, domain, preview_feature_key, send_event
+    )
+
+    connection.send_result(msg["id"])
+    send_event()

tests/components/labs/test_websocket_api.py

@@ -695,3 +695,199 @@ async def test_websocket_backup_timeout_handling(
 
     assert not msg["success"]
     assert msg["error"]["code"] == "unknown_error"
+
+
+async def test_websocket_subscribe_feature(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test subscribing to a specific preview feature."""
+    hass.config.components.add("kitchen_sink")
+    assert await async_setup(hass, {})
+    await hass.async_block_till_done()
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id(
+        {
+            "type": "labs/subscribe",
+            "domain": "kitchen_sink",
+            "preview_feature": "special_repair",
+        }
+    )
+    msg = await client.receive_json()
+
+    assert msg["success"]
+    assert msg["result"] is None
+
+    # Initial state is sent as event
+    event_msg = await client.receive_json()
+    assert event_msg["type"] == "event"
+    assert event_msg["event"] == {
+        "preview_feature": "special_repair",
+        "domain": "kitchen_sink",
+        "enabled": False,
+        "is_built_in": True,
+        "feedback_url": ANY,
+        "learn_more_url": ANY,
+        "report_issue_url": ANY,
+    }
+
+
+async def test_websocket_subscribe_feature_receives_updates(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test that subscription receives updates when feature is toggled."""
+    hass.config.components.add("kitchen_sink")
+    assert await async_setup(hass, {})
+    await hass.async_block_till_done()
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id(
+        {
+            "type": "labs/subscribe",
+            "domain": "kitchen_sink",
+            "preview_feature": "special_repair",
+        }
+    )
+    subscribe_msg = await client.receive_json()
+    assert subscribe_msg["success"]
+    subscription_id = subscribe_msg["id"]
+
+    # Initial state event
+    initial_event_msg = await client.receive_json()
+    assert initial_event_msg["id"] == subscription_id
+    assert initial_event_msg["type"] == "event"
+    assert initial_event_msg["event"]["enabled"] is False
+
+    await client.send_json_auto_id(
+        {
+            "type": "labs/update",
+            "domain": "kitchen_sink",
+            "preview_feature": "special_repair",
+            "enabled": True,
+        }
+    )
+
+    # Update event arrives before the update result
+    event_msg = await client.receive_json()
+    assert event_msg["id"] == subscription_id
+    assert event_msg["type"] == "event"
+    assert event_msg["event"] == {
+        "preview_feature": "special_repair",
+        "domain": "kitchen_sink",
+        "enabled": True,
+        "is_built_in": True,
+        "feedback_url": ANY,
+        "learn_more_url": ANY,
+        "report_issue_url": ANY,
+    }
+
+    update_msg = await client.receive_json()
+    assert update_msg["success"]
+
+
+async def test_websocket_subscribe_nonexistent_feature(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test subscribing to a preview feature that doesn't exist."""
+    assert await async_setup(hass, {})
+    await hass.async_block_till_done()
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id(
+        {
+            "type": "labs/subscribe",
+            "domain": "nonexistent",
+            "preview_feature": "feature",
+        }
+    )
+    msg = await client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == "not_found"
+    assert "not found" in msg["error"]["message"].lower()
+
+
+async def test_websocket_subscribe_does_not_require_admin(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    hass_admin_user: MockUser,
+) -> None:
+    """Test that subscribe does not require admin privileges."""
+    hass_admin_user.groups = []
+
+    hass.config.components.add("kitchen_sink")
+    assert await async_setup(hass, {})
+    await hass.async_block_till_done()
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id(
+        {
+            "type": "labs/subscribe",
+            "domain": "kitchen_sink",
+            "preview_feature": "special_repair",
+        }
+    )
+    msg = await client.receive_json()
+
+    assert msg["success"]
+
+    # Consume initial state event
+    await client.receive_json()
+
+
+async def test_websocket_subscribe_only_receives_subscribed_feature_updates(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test that subscription only receives updates for the subscribed feature."""
+    hass.config.components.add("kitchen_sink")
+    assert await async_setup(hass, {})
+    await hass.async_block_till_done()
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id(
+        {
+            "type": "labs/subscribe",
+            "domain": "kitchen_sink",
+            "preview_feature": "special_repair",
+        }
+    )
+    subscribe_msg = await client.receive_json()
+    assert subscribe_msg["success"]
+
+    # Consume initial state event
+    await client.receive_json()
+
+    # Fire an event for a different feature
+    hass.bus.async_fire(
+        EVENT_LABS_UPDATED,
+        {"domain": "other_domain", "preview_feature": "other_feature", "enabled": True},
+    )
+    await hass.async_block_till_done()
+
+    await client.send_json_auto_id(
+        {
+            "type": "labs/update",
+            "domain": "kitchen_sink",
+            "preview_feature": "special_repair",
+            "enabled": True,
+        }
+    )
+
+    # Event message arrives before the update result
+    # Should only receive event for subscribed feature, not the other one
+    event_msg = await client.receive_json()
+    assert event_msg["type"] == "event"
+    assert event_msg["event"]["domain"] == "kitchen_sink"
+    assert event_msg["event"]["preview_feature"] == "special_repair"
+
+    update_msg = await client.receive_json()
+    assert update_msg["success"]