Move websocket registration to async_setup in HTML5 integration (#181859)

homeassistant/components/html5/__init__.py

@@ -8,6 +8,7 @@ from homeassistant.helpers.typing import ConfigType
 
 from .const import DOMAIN
 from .services import async_setup_services
+from .websocket_api import async_register_websocket_api
 
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 
@@ -18,6 +19,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the HTML5 services."""
 
     async_setup_services(hass)
+    async_register_websocket_api(hass)
     return True
 
 

homeassistant/components/html5/notify.py

@@ -18,7 +18,6 @@ import probatio
 from py_vapid import Vapid
 from pywebpush import WebPusher, WebPushException, webpush_async
 
-from homeassistant.components import websocket_api
 from homeassistant.components.notify import (
     ATTR_DATA,
     ATTR_TARGET,
@@ -28,10 +27,9 @@ from homeassistant.components.notify import (
     NotifyEntity,
     NotifyEntityFeature,
 )
-from homeassistant.components.websocket_api import ActiveConnection
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import URL_ROOT
-from homeassistant.core import HomeAssistant, ServiceCall, callback
+from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -48,7 +46,6 @@ from .const import (
     ATTR_TTL,
     ATTR_VAPID_EMAIL,
     ATTR_VAPID_PRV_KEY,
-    ATTR_VAPID_PUB_KEY,
     DOMAIN,
     REGISTRATIONS_FILE,
     SERVICE_DISMISS,
@@ -71,11 +68,6 @@ DEFAULT_ICON = "/static/icons/favicon-192x192.png"
 
 ATTR_JWT = "jwt"
 
-WS_TYPE_APPKEY = "notify/html5/appkey"
-SCHEMA_WS_APPKEY = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend(
-    {probatio.Required("type"): WS_TYPE_APPKEY}
-)
-
 # The number of days after the moment a notification is sent that a JWT
 # is valid.
 JWT_VALID_DAYS = 7
@@ -123,20 +115,9 @@ async def async_get_service(
 
     registrations = await hass.async_add_executor_job(_load_config, json_path)
 
-    vapid_pub_key: str = discovery_info[ATTR_VAPID_PUB_KEY]
     vapid_prv_key: str = discovery_info[ATTR_VAPID_PRV_KEY]
     vapid_email: str = discovery_info[ATTR_VAPID_EMAIL]
 
-    @callback
-    def websocket_appkey(
-        _hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
-    ) -> None:
-        connection.send_message(websocket_api.result_message(msg["id"], vapid_pub_key))
-
-    websocket_api.async_register_command(
-        hass, WS_TYPE_APPKEY, websocket_appkey, SCHEMA_WS_APPKEY
-    )
-
     async_register_http_views(hass, json_path, registrations)
 
     session = async_get_clientsession(hass)

homeassistant/components/html5/websocket_api.py

@@ -0,0 +1,38 @@
+"""HTML5 Websocket API."""
+
+from typing import Any
+
+import probatio
+
+from homeassistant.components import websocket_api
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.service import async_get_config_entry
+
+from .const import ATTR_VAPID_PUB_KEY, DOMAIN
+
+WS_TYPE_APPKEY = "notify/html5/appkey"
+
+
+@callback
+def async_register_websocket_api(hass: HomeAssistant) -> None:
+    """Register the websocket API."""
+
+    websocket_api.async_register_command(hass, websocket_appkey)
+
+
+@websocket_api.websocket_command(
+    {
+        probatio.Required("type"): WS_TYPE_APPKEY,
+    }
+)
+@websocket_api.async_response
+async def websocket_appkey(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
+    """Handle request for the VAPID public key."""
+    entry = async_get_config_entry(hass, DOMAIN, None)
+    connection.send_message(
+        websocket_api.result_message(msg["id"], entry.data[ATTR_VAPID_PUB_KEY])
+    )

tests/components/html5/test_websocket_api.py

@@ -0,0 +1,32 @@
+"""Tests for HTML5 Websocket API."""
+
+from homeassistant.components.html5.websocket_api import WS_TYPE_APPKEY
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.core import HomeAssistant
+
+from .conftest import MOCK_CONF_PUB_KEY
+
+from tests.common import MockConfigEntry
+from tests.typing import WebSocketGenerator
+
+
+async def test_websocket_appkey(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test websocket appkey command."""
+
+    config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id({"type": WS_TYPE_APPKEY})
+    response = await client.receive_json()
+
+    assert response["success"]
+    assert response["result"] == MOCK_CONF_PUB_KEY