Add WS command frontend/subscribe_extra_js (#119833)

Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/frontend/__init__.py

@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
-from functools import lru_cache
+from functools import lru_cache, partial
 import logging
 import os
 import pathlib
@@ -33,6 +33,7 @@ from homeassistant.helpers.storage import Store
 from homeassistant.helpers.translation import async_get_translations
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import async_get_integration, bind_hass
+from homeassistant.util.hass_dict import HassKey
 
 from .storage import async_setup_frontend_storage
 
@@ -56,6 +57,10 @@ DATA_JS_VERSION = "frontend_js_version"
 DATA_EXTRA_MODULE_URL = "frontend_extra_module_url"
 DATA_EXTRA_JS_URL_ES5 = "frontend_extra_js_url_es5"
 
+DATA_WS_SUBSCRIBERS: HassKey[set[tuple[websocket_api.ActiveConnection, int]]] = HassKey(
+    "frontend_ws_subscribers"
+)
+
 THEMES_STORAGE_KEY = f"{DOMAIN}_theme"
 THEMES_STORAGE_VERSION = 1
 THEMES_SAVE_DELAY = 60
@@ -204,17 +209,24 @@ class UrlManager:
     on hass.data
     """
 
-    def __init__(self, urls: list[str]) -> None:
+    def __init__(
+        self,
+        on_change: Callable[[str, str], None],
+        urls: list[str],
+    ) -> None:
         """Init the url manager."""
+        self._on_change = on_change
         self.urls = frozenset(urls)
 
     def add(self, url: str) -> None:
         """Add a url to the set."""
         self.urls = frozenset([*self.urls, url])
+        self._on_change("added", url)
 
     def remove(self, url: str) -> None:
         """Remove a url from the set."""
         self.urls = self.urls - {url}
+        self._on_change("removed", url)
 
 
 class Panel:
@@ -363,6 +375,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     websocket_api.async_register_command(hass, websocket_get_themes)
     websocket_api.async_register_command(hass, websocket_get_translations)
     websocket_api.async_register_command(hass, websocket_get_version)
+    websocket_api.async_register_command(hass, websocket_subscribe_extra_js)
     hass.http.register_view(ManifestJSONView())
 
     conf = config.get(DOMAIN, {})
@@ -420,8 +433,27 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         sidebar_icon="hass:hammer",
     )
 
-    hass.data[DATA_EXTRA_MODULE_URL] = UrlManager(conf.get(CONF_EXTRA_MODULE_URL, []))
+    @callback
-    hass.data[DATA_EXTRA_JS_URL_ES5] = UrlManager(conf.get(CONF_EXTRA_JS_URL_ES5, []))
+    def async_change_listener(
+        resource_type: str,
+        change_type: str,
+        url: str,
+    ) -> None:
+        subscribers = hass.data[DATA_WS_SUBSCRIBERS]
+        json_msg = {
+            "change_type": change_type,
+            "item": {"type": resource_type, "url": url},
+        }
+        for connection, msg_id in subscribers:
+            connection.send_message(websocket_api.event_message(msg_id, json_msg))
+
+    hass.data[DATA_EXTRA_MODULE_URL] = UrlManager(
+        partial(async_change_listener, "module"), conf.get(CONF_EXTRA_MODULE_URL, [])
+    )
+    hass.data[DATA_EXTRA_JS_URL_ES5] = UrlManager(
+        partial(async_change_listener, "es5"), conf.get(CONF_EXTRA_JS_URL_ES5, [])
+    )
+    hass.data[DATA_WS_SUBSCRIBERS] = set()
 
     await _async_setup_themes(hass, conf.get(CONF_THEMES))
 
@@ -783,6 +815,24 @@ async def websocket_get_version(
         connection.send_result(msg["id"], {"version": frontend})
 
 
+@callback
+@websocket_api.websocket_command({"type": "frontend/subscribe_extra_js"})
+def websocket_subscribe_extra_js(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Subscribe to URL manager updates."""
+
+    subscribers = hass.data[DATA_WS_SUBSCRIBERS]
+    subscribers.add((connection, msg["id"]))
+
+    @callback
+    def cancel_subscription() -> None:
+        subscribers.remove((connection, msg["id"]))
+
+    connection.subscriptions[msg["id"]] = cancel_subscription
+    connection.send_message(websocket_api.result_message(msg["id"]))
+
+
 class PanelRespons(TypedDict):
     """Represent the panel response type."""
 

tests/components/frontend/test_init.py

@@ -409,7 +409,11 @@ async def test_missing_themes(hass: HomeAssistant, ws_client) -> None:
 
 
 @pytest.mark.usefixtures("mock_onboarded")
-async def test_extra_js(hass: HomeAssistant, mock_http_client_with_extra_js) -> None:
+async def test_extra_js(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    mock_http_client_with_extra_js,
+) -> None:
     """Test that extra javascript is loaded."""
 
     async def get_response():
@@ -423,6 +427,13 @@ async def test_extra_js(hass: HomeAssistant, mock_http_client_with_extra_js) ->
     assert '"/local/my_module.js"' in text
     assert '"/local/my_es5.js"' in text
 
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id({"type": "frontend/subscribe_extra_js"})
+    msg = await client.receive_json()
+
+    assert msg["success"] is True
+    subscription_id = msg["id"]
+
     # Test dynamically adding and removing extra javascript
     add_extra_js_url(hass, "/local/my_module_2.js", False)
     add_extra_js_url(hass, "/local/my_es5_2.js", True)
@@ -430,12 +441,38 @@ async def test_extra_js(hass: HomeAssistant, mock_http_client_with_extra_js) ->
     assert '"/local/my_module_2.js"' in text
     assert '"/local/my_es5_2.js"' in text
 
+    msg = await client.receive_json()
+    assert msg["id"] == subscription_id
+    assert msg["event"] == {
+        "change_type": "added",
+        "item": {"type": "module", "url": "/local/my_module_2.js"},
+    }
+    msg = await client.receive_json()
+    assert msg["id"] == subscription_id
+    assert msg["event"] == {
+        "change_type": "added",
+        "item": {"type": "es5", "url": "/local/my_es5_2.js"},
+    }
+
     remove_extra_js_url(hass, "/local/my_module_2.js", False)
     remove_extra_js_url(hass, "/local/my_es5_2.js", True)
     text = await get_response()
     assert '"/local/my_module_2.js"' not in text
     assert '"/local/my_es5_2.js"' not in text
 
+    msg = await client.receive_json()
+    assert msg["id"] == subscription_id
+    assert msg["event"] == {
+        "change_type": "removed",
+        "item": {"type": "module", "url": "/local/my_module_2.js"},
+    }
+    msg = await client.receive_json()
+    assert msg["id"] == subscription_id
+    assert msg["event"] == {
+        "change_type": "removed",
+        "item": {"type": "es5", "url": "/local/my_es5_2.js"},
+    }
+
     # Remove again should not raise
     remove_extra_js_url(hass, "/local/my_module_2.js", False)
     remove_extra_js_url(hass, "/local/my_es5_2.js", True)