Add websocket handlers to hassio (#46571)

homeassistant/components/hassio/__init__.py

@@ -10,7 +10,6 @@ from homeassistant.auth.const import GROUP_ID_ADMIN
 from homeassistant.components.homeassistant import SERVICE_CHECK_CONFIG
 import homeassistant.config as conf_util
 from homeassistant.const import (
-    ATTR_NAME,
     EVENT_CORE_CONFIG_UPDATE,
     SERVICE_HOMEASSISTANT_RESTART,
     SERVICE_HOMEASSISTANT_STOP,
@@ -24,15 +23,27 @@ from homeassistant.util.dt import utcnow
 
 from .addon_panel import async_setup_addon_panel
 from .auth import async_setup_auth_view
-from .const import ATTR_DISCOVERY
+from .const import (
+    ATTR_ADDON,
+    ATTR_ADDONS,
+    ATTR_DISCOVERY,
+    ATTR_FOLDERS,
+    ATTR_HOMEASSISTANT,
+    ATTR_INPUT,
+    ATTR_NAME,
+    ATTR_PASSWORD,
+    ATTR_SNAPSHOT,
+    DOMAIN,
+)
 from .discovery import async_setup_discovery_view
 from .handler import HassIO, HassioAPIError, api_data
 from .http import HassIOView
 from .ingress import async_setup_ingress_view
+from .websocket_api import async_load_websocket_api
 
 _LOGGER = logging.getLogger(__name__)
 
-DOMAIN = "hassio"
+
 STORAGE_KEY = DOMAIN
 STORAGE_VERSION = 1
 
@@ -62,13 +73,6 @@ SERVICE_SNAPSHOT_PARTIAL = "snapshot_partial"
 SERVICE_RESTORE_FULL = "restore_full"
 SERVICE_RESTORE_PARTIAL = "restore_partial"
 
-ATTR_ADDON = "addon"
-ATTR_INPUT = "input"
-ATTR_SNAPSHOT = "snapshot"
-ATTR_ADDONS = "addons"
-ATTR_FOLDERS = "folders"
-ATTR_HOMEASSISTANT = "homeassistant"
-ATTR_PASSWORD = "password"
 
 SCHEMA_NO_DATA = vol.Schema({})
 
@@ -101,6 +105,7 @@ SCHEMA_RESTORE_PARTIAL = SCHEMA_RESTORE_FULL.extend(
     }
 )
 
+
 MAP_SERVICE_API = {
     SERVICE_ADDON_START: ("/addons/{addon}/start", SCHEMA_ADDON, 60, False),
     SERVICE_ADDON_STOP: ("/addons/{addon}/stop", SCHEMA_ADDON, 60, False),
@@ -290,6 +295,8 @@ async def async_setup(hass, config):
         _LOGGER.error("Missing %s environment variable", env)
         return False
 
+    async_load_websocket_api(hass)
+
     host = os.environ["HASSIO"]
     websession = hass.helpers.aiohttp_client.async_get_clientsession()
     hass.data[DOMAIN] = hassio = HassIO(hass.loop, websession, host)

homeassistant/components/hassio/const.py

@@ -1,21 +1,42 @@
 """Hass.io const variables."""
 
-ATTR_ADDONS = "addons"
-ATTR_DISCOVERY = "discovery"
+DOMAIN = "hassio"
+
 ATTR_ADDON = "addon"
-ATTR_NAME = "name"
-ATTR_SERVICE = "service"
-ATTR_CONFIG = "config"
-ATTR_UUID = "uuid"
-ATTR_USERNAME = "username"
-ATTR_PASSWORD = "password"
-ATTR_PANELS = "panels"
-ATTR_ENABLE = "enable"
-ATTR_TITLE = "title"
-ATTR_ICON = "icon"
+ATTR_ADDONS = "addons"
 ATTR_ADMIN = "admin"
+ATTR_CONFIG = "config"
+ATTR_DATA = "data"
+ATTR_DISCOVERY = "discovery"
+ATTR_ENABLE = "enable"
+ATTR_FOLDERS = "folders"
+ATTR_HOMEASSISTANT = "homeassistant"
+ATTR_ICON = "icon"
+ATTR_INPUT = "input"
+ATTR_NAME = "name"
+ATTR_PANELS = "panels"
+ATTR_PASSWORD = "password"
+ATTR_SERVICE = "service"
+ATTR_SNAPSHOT = "snapshot"
+ATTR_TITLE = "title"
+ATTR_USERNAME = "username"
+ATTR_UUID = "uuid"
+ATTR_WS_EVENT = "event"
+ATTR_ENDPOINT = "endpoint"
+ATTR_METHOD = "method"
+ATTR_TIMEOUT = "timeout"
+
 
 X_HASSIO = "X-Hassio-Key"
 X_INGRESS_PATH = "X-Ingress-Path"
 X_HASS_USER_ID = "X-Hass-User-ID"
 X_HASS_IS_ADMIN = "X-Hass-Is-Admin"
+
+
+WS_TYPE = "type"
+WS_ID = "id"
+
+WS_TYPE_EVENT = "supervisor/event"
+WS_TYPE_API = "supervisor/api"
+
+EVENT_SUPERVISOR_EVENT = "supervisor_event"

homeassistant/components/hassio/websocket_api.py

@@ -0,0 +1,84 @@
+"""Websocekt API handlers for the hassio integration."""
+import logging
+
+import voluptuous as vol
+
+from homeassistant.components import websocket_api
+from homeassistant.components.websocket_api.connection import ActiveConnection
+from homeassistant.core import HomeAssistant, callback
+import homeassistant.helpers.config_validation as cv
+
+from .const import (
+    ATTR_DATA,
+    ATTR_ENDPOINT,
+    ATTR_METHOD,
+    ATTR_TIMEOUT,
+    ATTR_WS_EVENT,
+    DOMAIN,
+    EVENT_SUPERVISOR_EVENT,
+    WS_ID,
+    WS_TYPE,
+    WS_TYPE_API,
+    WS_TYPE_EVENT,
+)
+from .handler import HassIO
+
+SCHEMA_WEBSOCKET_EVENT = vol.Schema(
+    {vol.Required(ATTR_WS_EVENT): cv.string},
+    extra=vol.ALLOW_EXTRA,
+)
+
+_LOGGER: logging.Logger = logging.getLogger(__package__)
+
+
+@callback
+def async_load_websocket_api(hass: HomeAssistant):
+    """Set up the websocket API."""
+    websocket_api.async_register_command(hass, websocket_supervisor_event)
+    websocket_api.async_register_command(hass, websocket_supervisor_api)
+
+
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(WS_TYPE): WS_TYPE_EVENT,
+        vol.Required(ATTR_DATA): SCHEMA_WEBSOCKET_EVENT,
+    }
+)
+async def websocket_supervisor_event(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+):
+    """Publish events from the Supervisor."""
+    hass.bus.async_fire(EVENT_SUPERVISOR_EVENT, msg[ATTR_DATA])
+    connection.send_result(msg[WS_ID])
+
+
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(WS_TYPE): WS_TYPE_API,
+        vol.Required(ATTR_ENDPOINT): cv.string,
+        vol.Required(ATTR_METHOD): cv.string,
+        vol.Optional(ATTR_DATA): dict,
+        vol.Optional(ATTR_TIMEOUT): cv.string,
+    }
+)
+async def websocket_supervisor_api(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict
+):
+    """Websocket handler to call Supervisor API."""
+    supervisor: HassIO = hass.data[DOMAIN]
+    result = False
+    try:
+        result = await supervisor.send_command(
+            msg[ATTR_ENDPOINT],
+            method=msg[ATTR_METHOD],
+            timeout=msg.get(ATTR_TIMEOUT, 10),
+            payload=msg.get(ATTR_DATA, {}),
+        )
+    except hass.components.hassio.HassioAPIError as err:
+        _LOGGER.error("Failed to to call %s - %s", msg[ATTR_ENDPOINT], err)
+        connection.send_error(msg[WS_ID], err)
+    else:
+        connection.send_result(msg[WS_ID], result[ATTR_DATA])

tests/components/hassio/test_init.py

@@ -7,8 +7,21 @@ import pytest
 from homeassistant.auth.const import GROUP_ID_ADMIN
 from homeassistant.components import frontend
 from homeassistant.components.hassio import STORAGE_KEY
+from homeassistant.components.hassio.const import (
+    ATTR_DATA,
+    ATTR_ENDPOINT,
+    ATTR_METHOD,
+    EVENT_SUPERVISOR_EVENT,
+    WS_ID,
+    WS_TYPE,
+    WS_TYPE_API,
+    WS_TYPE_EVENT,
+)
+from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
+from tests.common import async_capture_events
+
 MOCK_ENVIRON = {"HASSIO": "127.0.0.1", "HASSIO_TOKEN": "abcdefgh"}
 
 
@@ -346,3 +359,58 @@ async def test_service_calls_core(hassio_env, hass, aioclient_mock):
         assert mock_check_config.called
 
     assert aioclient_mock.call_count == 5
+
+
+async def test_websocket_supervisor_event(
+    hassio_env, hass: HomeAssistant, hass_ws_client
+):
+    """Test Supervisor websocket event."""
+    assert await async_setup_component(hass, "hassio", {})
+    websocket_client = await hass_ws_client(hass)
+
+    test_event = async_capture_events(hass, EVENT_SUPERVISOR_EVENT)
+
+    await websocket_client.send_json(
+        {WS_ID: 1, WS_TYPE: WS_TYPE_EVENT, ATTR_DATA: {"event": "test"}}
+    )
+
+    assert await websocket_client.receive_json()
+    await hass.async_block_till_done()
+
+    assert test_event[0].data == {"event": "test"}
+
+
+async def test_websocket_supervisor_api(
+    hassio_env, hass: HomeAssistant, hass_ws_client, aioclient_mock
+):
+    """Test Supervisor websocket api."""
+    assert await async_setup_component(hass, "hassio", {})
+    websocket_client = await hass_ws_client(hass)
+    aioclient_mock.post(
+        "http://127.0.0.1/snapshots/new/partial",
+        json={"result": "ok", "data": {"slug": "sn_slug"}},
+    )
+
+    await websocket_client.send_json(
+        {
+            WS_ID: 1,
+            WS_TYPE: WS_TYPE_API,
+            ATTR_ENDPOINT: "/snapshots/new/partial",
+            ATTR_METHOD: "post",
+        }
+    )
+
+    msg = await websocket_client.receive_json()
+    assert msg["result"]["slug"] == "sn_slug"
+
+    await websocket_client.send_json(
+        {
+            WS_ID: 2,
+            WS_TYPE: WS_TYPE_API,
+            ATTR_ENDPOINT: "/supervisor/info",
+            ATTR_METHOD: "get",
+        }
+    )
+
+    msg = await websocket_client.receive_json()
+    assert msg["result"]["version_latest"] == "1.0.0"