Add a websocket API to the integration for event retain (#295)

* Add a new websocket API to the integration.

* Pass in the data parameter.

custom_components/frigate/__init__.py

@@ -56,6 +56,7 @@
     VodProxyView,
     VodSegmentProxyView,
 )
+from .ws_api import async_setup as ws_api_async_setup
 
 SCAN_INTERVAL = timedelta(seconds=5)
 
@@ -169,6 +170,8 @@ async def async_setup(hass: HomeAssistant, config: Config) -> bool:
 
     hass.data.setdefault(DOMAIN, {})
 
+    ws_api_async_setup(hass)
+
     session = async_get_clientsession(hass)
     hass.http.register_view(JSMPEGProxyView(session))
     hass.http.register_view(NotificationsProxyView(session))

custom_components/frigate/api.py

@@ -4,7 +4,7 @@
 import asyncio
 import logging
 import socket
-from typing import Any, Dict, List, cast
+from typing import Any, cast
 
 import aiohttp
 import async_timeout
@@ -47,7 +47,7 @@ async def async_get_version(self) -> str:
     async def async_get_stats(self) -> dict[str, Any]:
         """Get data from the API."""
         return cast(
-            Dict[str, Any],
+            dict[str, Any],
             await self.api_wrapper("get", str(URL(self._host) / "api/stats")),
         )
 
@@ -75,7 +75,7 @@ async def async_get_events(
         }
 
         return cast(
-            List[Dict[str, Any]],
+            list[dict[str, Any]],
             await self.api_wrapper(
                 "get",
                 str(
@@ -98,7 +98,7 @@ async def async_get_event_summary(
         }
 
         return cast(
-            List[Dict[str, Any]],
+            list[dict[str, Any]],
             await self.api_wrapper(
                 "get",
                 str(
@@ -112,14 +112,25 @@ async def async_get_event_summary(
     async def async_get_config(self) -> dict[str, Any]:
         """Get data from the API."""
         return cast(
-            Dict[str, Any],
+            dict[str, Any],
             await self.api_wrapper("get", str(URL(self._host) / "api/config")),
         )
 
     async def async_get_path(self, path: str) -> Any:
         """Get data from the API."""
         return await self.api_wrapper("get", str(URL(self._host) / f"{path}/"))
 
+    async def async_retain(self, event_id: str, retain: bool) -> dict[str, Any]:
+        """Un/Retain an event."""
+        return cast(
+            dict[str, Any],
+            await self.api_wrapper(
+                "post" if retain else "delete",
+                str(URL(self._host) / f"api/events/{event_id}/retain"),
+                decode_json=True,
+            ),
+        )
+
     async def api_wrapper(
         self,
         method: str,
@@ -136,23 +147,15 @@ async def api_wrapper(
 
         try:
             async with async_timeout.timeout(TIMEOUT):
-                if method == "get":
-                    response = await self._session.get(
-                        url, headers=headers, raise_for_status=True
+                func = getattr(self._session, method)
+                if func:
+                    response = await func(
+                        url, headers=headers, raise_for_status=True, json=data
                     )
                     if decode_json:
                         return await response.json()
                     return await response.text()
 
-                if method == "put":
-                    await self._session.put(url, headers=headers, json=data)
-
-                elif method == "patch":
-                    await self._session.patch(url, headers=headers, json=data)
-
-                elif method == "post":
-                    await self._session.post(url, headers=headers, json=data)
-
         except asyncio.TimeoutError as exc:
             _LOGGER.error(
                 "Timeout error fetching information from %s: %s",

custom_components/frigate/media_source.py

@@ -32,8 +32,9 @@
 
 from . import get_friendly_name
 from .api import FrigateApiClient, FrigateApiClientError
-from .const import ATTR_CLIENT, CONF_MEDIA_BROWSER_ENABLE, DOMAIN, NAME
+from .const import CONF_MEDIA_BROWSER_ENABLE, DOMAIN, NAME
 from .views import (
+    get_client_for_frigate_instance_id,
     get_config_entry_for_frigate_instance_id,
     get_default_config_entry,
     get_frigate_instance_id_for_config_entry,
@@ -565,16 +566,11 @@ def _is_allowed_as_media_source(self, instance_id: str) -> bool:
 
     def _get_client(self, identifier: Identifier) -> FrigateApiClient:
         """Get client for a given identifier."""
-        config_entry = get_config_entry_for_frigate_instance_id(
+        client = get_client_for_frigate_instance_id(
             self.hass, identifier.frigate_instance_id
         )
-
-        if config_entry:
-            client: FrigateApiClient = (
-                self.hass.data[DOMAIN].get(config_entry.entry_id, {}).get(ATTR_CLIENT)
-            )
-            if client:
-                return client
+        if client:
+            return client
 
         raise MediaSourceError(
             "Could not find client for frigate instance id: %s"

custom_components/frigate/views.py

@@ -14,7 +14,9 @@
 from multidict import CIMultiDict
 from yarl import URL
 
+from custom_components.frigate.api import FrigateApiClient
 from custom_components.frigate.const import (
+    ATTR_CLIENT,
     ATTR_CLIENT_ID,
     ATTR_CONFIG,
     ATTR_MQTT,
@@ -67,6 +69,20 @@ def get_config_entry_for_frigate_instance_id(
     return None
 
 
+def get_client_for_frigate_instance_id(
+    hass: HomeAssistant, frigate_instance_id: str
+) -> FrigateApiClient | None:
+    """Get a client for a given frigate_instance_id."""
+
+    config_entry = get_config_entry_for_frigate_instance_id(hass, frigate_instance_id)
+    if config_entry:
+        return cast(
+            FrigateApiClient,
+            hass.data[DOMAIN].get(config_entry.entry_id, {}).get(ATTR_CLIENT),
+        )
+    return None
+
+
 def get_frigate_instance_id_for_config_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,

custom_components/frigate/ws_api.py

@@ -0,0 +1,67 @@
+"""Frigate HTTP views."""
+from __future__ import annotations
+
+import logging
+
+import voluptuous as vol
+
+from custom_components.frigate.api import FrigateApiClient, FrigateApiClientError
+from custom_components.frigate.views import get_client_for_frigate_instance_id
+from homeassistant.components import websocket_api
+from homeassistant.core import HomeAssistant
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+def async_setup(hass: HomeAssistant) -> None:
+    """Set up the recorder websocket API."""
+    websocket_api.async_register_command(hass, ws_retain_event)
+
+
+def _get_client_or_send_error(
+    hass: HomeAssistant,
+    instance_id: str,
+    msg_id: int,
+    connection: websocket_api.ActiveConnection,
+) -> FrigateApiClient | None:
+    """Get the API client or send an error that it cannot be found."""
+    client = get_client_for_frigate_instance_id(hass, instance_id)
+    if client is None:
+        connection.send_error(
+            msg_id,
+            websocket_api.const.ERR_NOT_FOUND,
+            f"Unable to find Frigate instance with ID: {instance_id}",
+        )
+        return None
+    return client
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "frigate/event/retain",
+        vol.Required("instance_id"): str,
+        vol.Required("event_id"): str,
+        vol.Required("retain"): bool,
+    }
+)  # type: ignore[misc]
+@websocket_api.async_response  # type: ignore[misc]
+async def ws_retain_event(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict,
+) -> None:
+    """Un/Retain an event."""
+    client = _get_client_or_send_error(hass, msg["instance_id"], msg["id"], connection)
+    if not client:
+        return
+    try:
+        connection.send_result(
+            msg["id"], await client.async_retain(msg["event_id"], msg["retain"])
+        )
+    except FrigateApiClientError:
+        connection.send_error(
+            msg["id"],
+            "frigate_error",
+            f"API error whilst un/retaining event {msg['event_id']} "
+            f"for Frigate instance {msg['instance_id']}",
+        )

tests/test_api.py

@@ -264,3 +264,32 @@ async def version_handler(request: web.Request) -> web.Response:
 
     frigate_client = FrigateApiClient(str(server.make_url("/")), aiohttp_session)
     assert await frigate_client.async_get_version() == TEST_SERVER_VERSION
+
+
+async def test_async_retain(
+    aiohttp_session: aiohttp.ClientSession, aiohttp_server: Any
+) -> None:
+    """Test async_retain."""
+
+    post_success = {"success": True, "message": "Post success"}
+    post_handler = Mock(return_value=web.json_response(post_success))
+
+    delete_success = {"success": True, "message": "Delete success"}
+    delete_handler = Mock(return_value=web.json_response(delete_success))
+
+    event_id = "1656282822.206673-bovnfg"
+    server = await start_frigate_server(
+        aiohttp_server,
+        [
+            web.post(f"/api/events/{event_id}/retain", post_handler),
+            web.delete(f"/api/events/{event_id}/retain", delete_handler),
+        ],
+    )
+
+    frigate_client = FrigateApiClient(str(server.make_url("/")), aiohttp_session)
+    assert await frigate_client.async_retain(event_id, True) == post_success
+    assert post_handler.called
+    assert not delete_handler.called
+
+    assert await frigate_client.async_retain(event_id, False) == delete_success
+    assert delete_handler.called

tests/test_ws_api.py

@@ -0,0 +1,121 @@
+"""Test the Frigate HA websocket API."""
+from __future__ import annotations
+
+import logging
+from typing import Any
+from unittest.mock import AsyncMock
+
+from custom_components.frigate.api import FrigateApiClientError
+from homeassistant.core import HomeAssistant
+
+from tests import (
+    TEST_FRIGATE_INSTANCE_ID,
+    create_mock_frigate_client,
+    setup_mock_frigate_config_entry,
+)
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+TEST_EVENT_ID = "1656282822.206673-bovnfg"
+
+
+async def test_retain_success(hass: HomeAssistant, hass_ws_client: Any) -> None:
+    """Test un/retaining an event."""
+
+    mock_client = create_mock_frigate_client()
+    await setup_mock_frigate_config_entry(hass, client=mock_client)
+
+    ws_client = await hass_ws_client()
+    retain_json: dict[str, Any] = {
+        "id": 1,
+        "type": "frigate/event/retain",
+        "instance_id": TEST_FRIGATE_INSTANCE_ID,
+        "event_id": TEST_EVENT_ID,
+        "retain": True,
+    }
+
+    retain_success = {"retain": "success"}
+    mock_client.async_retain = AsyncMock(return_value=retain_success)
+    await ws_client.send_json(retain_json)
+
+    response = await ws_client.receive_json()
+    mock_client.async_retain.assert_called_with(TEST_EVENT_ID, True)
+    assert response["success"]
+    assert response["result"] == retain_success
+
+    unretain_success = {"unretain": "success"}
+    mock_client.async_retain = AsyncMock(return_value=unretain_success)
+    await ws_client.send_json(
+        {
+            **retain_json,
+            "id": 2,
+            "retain": False,
+        }
+    )
+
+    response = await ws_client.receive_json()
+    mock_client.async_retain.assert_called_with(TEST_EVENT_ID, False)
+    assert response["success"]
+    assert response["result"] == unretain_success
+
+
+async def test_retain_missing_args(hass: HomeAssistant, hass_ws_client: Any) -> None:
+    """Test retaining an event with missing arguments."""
+
+    await setup_mock_frigate_config_entry(hass)
+
+    ws_client = await hass_ws_client()
+    retain_json = {
+        "id": 1,
+        "type": "frigate/event/retain",
+    }
+
+    await ws_client.send_json(retain_json)
+    response = await ws_client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == "invalid_format"
+
+
+async def test_retain_instance_not_found(
+    hass: HomeAssistant, hass_ws_client: Any
+) -> None:
+    """Test retaining an event with an instance that is not found."""
+
+    await setup_mock_frigate_config_entry(hass)
+
+    ws_client = await hass_ws_client()
+    retain_json = {
+        "id": 1,
+        "type": "frigate/event/retain",
+        "instance_id": "THIS-IS-NOT-A-REAL-INSTANCE-ID",
+        "event_id": TEST_EVENT_ID,
+        "retain": True,
+    }
+
+    await ws_client.send_json(retain_json)
+    response = await ws_client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == "not_found"
+
+
+async def test_retain_api_error(hass: HomeAssistant, hass_ws_client: Any) -> None:
+    """Test retaining an event when the API has an error."""
+
+    mock_client = create_mock_frigate_client()
+    await setup_mock_frigate_config_entry(hass, client=mock_client)
+
+    ws_client = await hass_ws_client()
+    retain_json = {
+        "id": 1,
+        "type": "frigate/event/retain",
+        "instance_id": TEST_FRIGATE_INSTANCE_ID,
+        "event_id": TEST_EVENT_ID,
+        "retain": True,
+    }
+
+    mock_client.async_retain = AsyncMock(side_effect=FrigateApiClientError)
+
+    await ws_client.send_json(retain_json)
+    response = await ws_client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == "frigate_error"