Add subscribe preview feature endpoint to labs (home-assistant#157976)

Author: piitaya

homeassistant/components/labs/__init__.py

@@ -7,18 +7,18 @@
 
 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
 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.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()