Add ice servers to async WebRTC offer (#4)

* Add ice servers to async WebRTC offer

* Update to reflect upstream changes

* Fix code and test

* Guard against unexpected messages

* Fix typing

* Increase coverage

Author: edenhaus

.pre-commit-config.yaml

@@ -7,6 +7,7 @@ repos:
         name: 🐶 Ruff lint
         args:
           - --fix
+          # - --unsafe-fixes
 
       - id: ruff-format
         name: 🐶 Ruff format

go2rtc_client/ws/client.py

@@ -10,7 +10,7 @@
 
 from go2rtc_client.exceptions import handle_error
 
-from .messages import BaseMessage
+from .messages import BaseMessage, ReceiveMessages, SendMessages, WebRTC, WsMessage
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -43,7 +43,7 @@ def __init__(
         self._params = params
         self._client: ClientWebSocketResponse | None = None
         self._rx_task: asyncio.Task[None] | None = None
-        self._subscribers: list[Callable[[BaseMessage], None]] = []
+        self._subscribers: list[Callable[[ReceiveMessages], None]] = []
         self._connect_lock = asyncio.Lock()
 
     @property
@@ -77,7 +77,7 @@ async def close(self) -> None:
             await client.close()
 
     @handle_error
-    async def send(self, message: BaseMessage) -> None:
+    async def send(self, message: SendMessages) -> None:
         """Send a message."""
         if not self.connected:
             await self.connect()
@@ -90,10 +90,15 @@ async def send(self, message: BaseMessage) -> None:
     def _process_text_message(self, data: Any) -> None:
         """Process text message."""
         try:
-            message = BaseMessage.from_json(data)
+            message: WsMessage = BaseMessage.from_json(data)
         except Exception:  # pylint: disable=broad-except
             _LOGGER.exception("Invalid message received: %s", data)
         else:
+            if isinstance(message, WebRTC):
+                message = message.value
+            if not isinstance(message, ReceiveMessages):
+                _LOGGER.error("Received unexpected message: %s", message)
+                return
             for subscriber in self._subscribers:
                 try:
                     subscriber(message)
@@ -134,7 +139,9 @@ async def _receive_messages(self) -> None:
             if self.connected:
                 await self.close()
 
-    def subscribe(self, callback: Callable[[BaseMessage], None]) -> Callable[[], None]:
+    def subscribe(
+        self, callback: Callable[[ReceiveMessages], None]
+    ) -> Callable[[], None]:
         """Subscribe to messages."""
 
         def _unsubscribe() -> None:

go2rtc_client/ws/messages.py

@@ -3,20 +3,34 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, ClassVar
+from typing import Annotated, Any, ClassVar
 
 from mashumaro import field_options
 from mashumaro.config import BaseConfig
 from mashumaro.mixins.orjson import DataClassORJSONMixin
 from mashumaro.types import Discriminator
+from webrtc_models import (
+    RTCIceServer,  # noqa: TCH002 # Mashumaro needs the import to generate the correct code
+)
 
 
 @dataclass(frozen=True)
-class BaseMessage(DataClassORJSONMixin):
-    """Base message class."""
+class WsMessage:
+    """Websocket message."""
 
     TYPE: ClassVar[str]
 
+    def __post_serialize__(self, d: dict[Any, Any]) -> dict[Any, Any]:
+        """Add type to serialized dict."""
+        # ClassVar will not serialize by default
+        d["type"] = self.TYPE
+        return d
+
+
+@dataclass(frozen=True)
+class BaseMessage(WsMessage, DataClassORJSONMixin):
+    """Base message class."""
+
     class Config(BaseConfig):
         """Config for BaseMessage."""
 
@@ -27,12 +41,6 @@ class Config(BaseConfig):
             variant_tagger_fn=lambda cls: cls.TYPE,
         )
 
-    def __post_serialize__(self, d: dict[Any, Any]) -> dict[Any, Any]:
-        """Add type to serialized dict."""
-        # ClassVar will not serialize by default
-        d["type"] = self.TYPE
-        return d
-
 
 @dataclass(frozen=True)
 class WebRTCCandidate(BaseMessage):
@@ -43,19 +51,55 @@ class WebRTCCandidate(BaseMessage):
 
 
 @dataclass(frozen=True)
-class WebRTCOffer(BaseMessage):
+class WebRTC(BaseMessage):
+    """WebRTC message."""
+
+    TYPE = "webrtc"
+    value: Annotated[
+        WebRTCOffer | WebRTCValue,
+        Discriminator(
+            field="type",
+            include_subtypes=True,
+            variant_tagger_fn=lambda cls: cls.TYPE,
+        ),
+    ]
+
+
+@dataclass(frozen=True)
+class WebRTCValue(WsMessage):
+    """WebRTC value for WebRTC message."""
+
+    sdp: str
+
+
+@dataclass(frozen=True)
+class WebRTCOffer(WebRTCValue):
     """WebRTC offer message."""
 
-    TYPE = "webrtc/offer"
-    offer: str = field(metadata=field_options(alias="value"))
+    TYPE = "offer"
+    ice_servers: list[RTCIceServer]
+
+    def __pre_serialize__(self) -> WebRTCOffer:
+        """Pre serialize.
+
+        Go2rtc supports only ice_servers with urls as list of strings.
+        """
+        for server in self.ice_servers:
+            if isinstance(server.urls, str):
+                server.urls = [server.urls]
+
+        return self
+
+    def to_json(self, **kwargs: Any) -> str:
+        """Convert to json."""
+        return WebRTC(self).to_json(**kwargs)
 
 
 @dataclass(frozen=True)
-class WebRTCAnswer(BaseMessage):
+class WebRTCAnswer(WebRTCValue):
     """WebRTC answer message."""
 
-    TYPE = "webrtc/answer"
-    answer: str = field(metadata=field_options(alias="value"))
+    TYPE = "answer"
 
 
 @dataclass(frozen=True)

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "awesomeversion>=24.6.0",
     "mashumaro~=3.13",
     "orjson>=3.10.7",
+    "webrtc-models>=0.1.0",
 ]
 version = "0.0.0"
 

tests/ws/test_client.py

@@ -17,11 +17,18 @@
 from aiohttp.web import WebSocketResponse
 from multidict import CIMultiDict, CIMultiDictProxy
 import pytest
+from webrtc_models import RTCIceServer
 from yarl import URL
 
 from go2rtc_client.exceptions import Go2RtcClientError
-from go2rtc_client.ws.client import Go2RtcWsClient
-from go2rtc_client.ws.messages import BaseMessage, WebRTCAnswer, WebRTCCandidate
+from go2rtc_client.ws import (
+    Go2RtcWsClient,
+    ReceiveMessages,
+    SendMessages,
+    WebRTCAnswer,
+    WebRTCCandidate,
+    WebRTCOffer,
+)
 
 
 class TestServer:
@@ -111,7 +118,27 @@ async def test_connect_parallel(server: TestServer) -> None:
         assert client.connected
 
 
-async def test_send(ws_client: Go2RtcWsClient, server: TestServer) -> None:
+@pytest.mark.parametrize(
+    ("message", "expected"),
+    [
+        (WebRTCCandidate("test"), '{"value":"test","type":"webrtc/candidate"}'),
+        (
+            WebRTCOffer("test", []),
+            '{"value":{"sdp":"test","ice_servers":[],"type":"offer"},"type":"webrtc"}',
+        ),
+        (
+            WebRTCOffer("test", [RTCIceServer("url")]),
+            '{"value":{"sdp":"test","ice_servers":[{"urls":["url"]}],"type":"offer"},"type":"webrtc"}',
+        ),
+        (
+            WebRTCOffer("test", [RTCIceServer(["url1", "url2"])]),
+            '{"value":{"sdp":"test","ice_servers":[{"urls":["url1","url2"]}],"type":"offer"},"type":"webrtc"}',
+        ),
+    ],
+)
+async def test_send(
+    ws_client: Go2RtcWsClient, server: TestServer, message: SendMessages, expected: str
+) -> None:
     """Test sending a message through the WebSocket."""
     received_message = None
 
@@ -121,28 +148,31 @@ def on_message(msg: WSMessage) -> None:
 
     server.on_message = on_message
 
-    await ws_client.send(WebRTCCandidate("test"))
+    await ws_client.send(message)
     await asyncio.sleep(0.1)
-    assert received_message == '{"value":"test","type":"webrtc/candidate"}'
+    assert received_message == expected
 
 
 @pytest.mark.parametrize(
     ("message", "expected"),
     [
         ('{"value":"test","type":"webrtc/candidate"}', WebRTCCandidate("test")),
-        ('{"value":"test","type":"webrtc/answer"}', WebRTCAnswer("test")),
+        (
+            '{"value":{"type":"answer", "sdp":"test"},"type":"webrtc"}',
+            WebRTCAnswer("test"),
+        ),
     ],
 )
 async def test_receive(
     ws_client_connected: Go2RtcWsClient,
     server: TestServer,
     message: str,
-    expected: BaseMessage,
+    expected: ReceiveMessages,
 ) -> None:
     """Test receiving a message through the WebSocket."""
     received_message = None
 
-    def on_message(message: BaseMessage) -> None:
+    def on_message(message: ReceiveMessages) -> None:
         nonlocal received_message
         received_message = message
 
@@ -230,7 +260,7 @@ async def test_subscribe_unsubscribe(ws_client: Go2RtcWsClient) -> None:
     # pylint: disable=protected-access
     assert ws_client._subscribers == []
 
-    def on_message(_: BaseMessage) -> None:
+    def on_message(_: ReceiveMessages) -> None:
         pass
 
     unsub = ws_client.subscribe(on_message)
@@ -249,14 +279,14 @@ async def test_subscriber_raised(
 ) -> None:
     """Test any exception raised by any subscriber will be handled."""
 
-    def on_message_raise(_: BaseMessage) -> None:
+    def on_message_raise(_: ReceiveMessages) -> None:
         raise ValueError
 
     ws_client_connected.subscribe(on_message_raise)
 
     received_message = None
 
-    def on_message(message: BaseMessage) -> None:
+    def on_message(message: ReceiveMessages) -> None:
         nonlocal received_message
         received_message = message
 
@@ -294,6 +324,18 @@ def on_message(message: BaseMessage) -> None:
             WSMessage(WSMsgType.ERROR, "error", None),
             ("go2rtc_client.ws.client", logging.ERROR, "Error received: error"),
         ),
+        (
+            WSMessage(
+                WSMsgType.TEXT,
+                '{"value":{"sdp":"test","ice_servers":[],"type":"offer"},"type":"webrtc"}',
+                None,
+            ),
+            (
+                "go2rtc_client.ws.client",
+                logging.ERROR,
+                "Received unexpected message: WebRTCOffer(sdp='test', ice_servers=[])",
+            ),
+        ),
     ],
 )
 async def test_unexpected_messages(