Add probing a stream (#23)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Author: edenhaus

go2rtc_client/models.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, Literal
+from typing import Generic, Literal, TypeVar
 
 from awesomeversion import AwesomeVersion
 from mashumaro import field_options
@@ -19,6 +19,20 @@ def deserialize(self, value: str) -> AwesomeVersion:
         return AwesomeVersion(value)
 
 
+_T = TypeVar("_T")
+
+
+class _EmptyListInsteadNoneSerializer(SerializationStrategy, Generic[_T]):
+    def serialize(self, value: list[_T]) -> list[_T]:
+        return value
+
+    def deserialize(self, value: list[_T] | None) -> list[_T]:
+        if value is None:
+            return []
+
+        return value
+
+
 @dataclass
 class ApplicationInfo(DataClassORJSONMixin):
     """Application info model.
@@ -39,26 +53,22 @@ class Streams(DataClassORJSONMixin):
 
 
 @dataclass
-class Stream:
+class Stream(DataClassORJSONMixin):
     """Stream model."""
 
-    producers: list[Producer]
-
-    @classmethod
-    def __pre_deserialize__(cls, d: dict[Any, Any]) -> dict[Any, Any]:
-        """Pre deserialize."""
-        # Ensure producers is always a list
-        if "producers" in d and d["producers"] is None:
-            d["producers"] = []
-
-        return d
+    producers: list[Producer] = field(
+        metadata=field_options(serialization_strategy=_EmptyListInsteadNoneSerializer())
+    )
 
 
 @dataclass
 class Producer:
     """Producer model."""
 
     url: str
+    medias: list[str] = field(
+        metadata=field_options(serialization_strategy=_EmptyListInsteadNoneSerializer())
+    )
 
 
 @dataclass

go2rtc_client/rest.py

@@ -133,6 +133,19 @@ async def add(self, name: str, source: str) -> None:
             params={"name": name, "src": [source, f"ffmpeg:{name}#audio=opus"]},
         )
 
+    @handle_error
+    async def probe(
+        self, stream_name: str, *, audio: str | None = None, video: str | None = None
+    ) -> Stream:
+        """Probe a stream."""
+        params = {"src": stream_name}
+        if audio:
+            params["audio"] = audio
+        if video:
+            params["video"] = video
+        resp = await self._client.request("GET", self.PATH, params=params)
+        return Stream.from_dict(await resp.json())
+
 
 class Go2RtcRestClient:
     """Rest client for go2rtc server."""

tests/__snapshots__/test_rest.ambr

@@ -9,6 +9,57 @@
     'version': '1.9.4',
   })
 # ---
+# name: test_probe_success[audio and video][deserialized]
+  dict({
+    'producers': list([
+      dict({
+        'medias': list([
+          'video, recvonly, H264',
+          'audio, recvonly, MPEG4-GENERIC/16000',
+          'audio, sendonly, PCMU/8000',
+        ]),
+        'url': 'rtsp://x:x@192.168.x.x:554/Preview_01_sub',
+      }),
+    ]),
+  })
+# ---
+# name: test_probe_success[audio and video][serialized]
+  '{"producers":[{"url":"rtsp://x:x@192.168.x.x:554/Preview_01_sub","medias":["video, recvonly, H264","audio, recvonly, MPEG4-GENERIC/16000","audio, sendonly, PCMU/8000"]}]}'
+# ---
+# name: test_probe_success[audio only][deserialized]
+  dict({
+    'producers': list([
+      dict({
+        'medias': list([
+          'video, recvonly, H264',
+          'audio, recvonly, MPEG4-GENERIC/16000',
+          'audio, sendonly, PCMU/8000',
+        ]),
+        'url': 'rtsp://x:x@192.168.x.x:554/Preview_01_sub',
+      }),
+    ]),
+  })
+# ---
+# name: test_probe_success[audio only][serialized]
+  '{"producers":[{"url":"rtsp://x:x@192.168.x.x:554/Preview_01_sub","medias":["video, recvonly, H264","audio, recvonly, MPEG4-GENERIC/16000","audio, sendonly, PCMU/8000"]}]}'
+# ---
+# name: test_probe_success[video only][deserialized]
+  dict({
+    'producers': list([
+      dict({
+        'medias': list([
+          'video, recvonly, H264',
+          'audio, recvonly, MPEG4-GENERIC/16000',
+          'audio, sendonly, PCMU/8000',
+        ]),
+        'url': 'rtsp://x:x@192.168.x.x:554/Preview_01_sub',
+      }),
+    ]),
+  })
+# ---
+# name: test_probe_success[video only][serialized]
+  '{"producers":[{"url":"rtsp://x:x@192.168.x.x:554/Preview_01_sub","medias":["video, recvonly, H264","audio, recvonly, MPEG4-GENERIC/16000","audio, sendonly, PCMU/8000"]}]}'
+# ---
 # name: test_streams_get[empty]
   dict({
   })
@@ -18,6 +69,10 @@
     'camera.12mp_fluent': dict({
       'producers': list([
         dict({
+          'medias': list([
+            'video, recvonly, H264',
+            'audio, recvonly, MPEG4-GENERIC/16000',
+          ]),
           'url': 'rtsp://test:test@192.168.10.105:554/Preview_06_sub',
         }),
       ]),

tests/fixtures/probe_camera_offline.txt

@@ -0,0 +1 @@
+streams: dial tcp 192.168.10.107:554: i/o timeout

tests/fixtures/probe_success.json

@@ -0,0 +1,95 @@
+{
+  "producers": [
+    {
+      "id": 2,
+      "format_name": "rtsp",
+      "protocol": "rtsp+tcp",
+      "remote_addr": "192.168.x.x:554",
+      "url": "rtsp://x:x@192.168.x.x:554/Preview_01_sub",
+      "sdp": "v=0\r\no=- 1730880192000275 1 IN IP4 192.168.x.x\r\ns=Session streamed by \"preview\"\r\nt=0 0\r\na=tool:BC Streaming Media v202210012022.10.01\r\na=type:broadcast\r\na=control:*\r\na=range:npt=now-\r\na=x-qt-text-nam:Session streamed by \"preview\"\r\nm=video 0 RTP/AVP 96\r\nc=IN IP4 0.0.0.0\r\nb=AS:8192\r\na=rtpmap:96 H264/90000\r\na=fmtp:96 packetization-mode=1;profile-level-id=640033;sprop-parameter-sets=Z2QAM6wVFKCgPZA=,aO48sA==\r\na=recvonly\r\na=control:track1\r\nm=audio 0 RTP/AVP 97\r\nc=IN IP4 0.0.0.0\r\nb=AS:8192\r\na=rtpmap:97 MPEG4-GENERIC/16000\r\na=fmtp:97 streamtype=5;profile-level-id=1;mode=AAC-hbr;sizelength=13;indexlength=3;indexdeltalength=3;config=1408\r\na=recvonly\r\na=control:track2\r\nm=audio 0 RTP/AVP 0\r\na=control:track3\r\na=rtpmap:0 PCMU/8000\r\na=sendonly",
+      "user_agent": "go2rtc/1.9.6",
+      "medias": [
+        "video, recvonly, H264",
+        "audio, recvonly, MPEG4-GENERIC/16000",
+        "audio, sendonly, PCMU/8000"
+      ],
+      "receivers": [
+        {
+          "id": 3,
+          "codec": {
+            "codec_name": "h264",
+            "codec_type": "video",
+            "level": 51,
+            "profile": "High"
+          },
+          "childs": [4]
+        },
+        {
+          "id": 5,
+          "codec": {
+            "codec_name": "aac",
+            "codec_type": "audio",
+            "sample_rate": 16000
+          },
+          "childs": [6]
+        }
+      ],
+      "senders": [
+        {
+          "id": 8,
+          "codec": {
+            "codec_name": "pcm_mulaw",
+            "codec_type": "audio",
+            "sample_rate": 8000
+          },
+          "parent": 7
+        }
+      ]
+    }
+  ],
+  "consumers": [
+    {
+      "id": 1,
+      "format_name": "probe",
+      "protocol": "http",
+      "remote_addr": "172.19.x.x:47622 forwarded 192.168.x.x",
+      "user_agent": "Mozilla/5.0 (X11; Linux x86_64; rv:132.0) Gecko/20100101 Firefox/132.0",
+      "medias": [
+        "video, sendonly, ALL",
+        "audio, sendonly, ALL",
+        "audio, recvonly, ANY"
+      ],
+      "receivers": [
+        {
+          "id": 7,
+          "codec": {
+            "codec_name": "ANY",
+            "codec_type": ""
+          },
+          "childs": [8]
+        }
+      ],
+      "senders": [
+        {
+          "id": 4,
+          "codec": {
+            "codec_name": "h264",
+            "codec_type": "video",
+            "level": 51,
+            "profile": "High"
+          },
+          "parent": 3
+        },
+        {
+          "id": 6,
+          "codec": {
+            "codec_name": "aac",
+            "codec_type": "audio",
+            "sample_rate": 16000
+          },
+          "parent": 5
+        }
+      ]
+    }
+  ]
+}

tests/test_rest.py

@@ -11,8 +11,8 @@
 from awesomeversion import AwesomeVersion
 import pytest
 
-from go2rtc_client.exceptions import Go2RtcVersionError
-from go2rtc_client.models import WebRTCSdpOffer
+from go2rtc_client.exceptions import Go2RtcClientError, Go2RtcVersionError
+from go2rtc_client.models import Stream, WebRTCSdpOffer
 from go2rtc_client.rest import _ApplicationClient, _StreamClient, _WebRTCClient
 from tests import load_fixture
 
@@ -146,3 +146,74 @@ async def test_webrtc_offer(
         WebRTCSdpOffer("v=0..."),
     )
     assert resp == snapshot
+
+
+async def _test_probe(
+    responses: aioresponses,
+    rest_client: Go2RtcRestClient,
+    filename: str,
+    status_code: int,
+    additional_params: dict[str, str],
+) -> Stream:
+    """Test probing a stream."""
+    camera = "camera.test"
+    params = [f"{k}={v}" for k, v in additional_params.items()]
+    responses.get(
+        f"{URL}{_StreamClient.PATH}?src={camera}&{'&'.join(params)}",
+        status=status_code,
+        body=load_fixture(filename),
+    )
+    return await rest_client.streams.probe(camera, **additional_params)
+
+
+@pytest.mark.parametrize(
+    "additional_params",
+    [
+        {"audio": "all", "video": "all"},
+        {"audio": "all"},
+        {"video": "all"},
+    ],
+    ids=[
+        "audio and video",
+        "audio only",
+        "video only",
+    ],
+)
+async def test_probe_success(
+    responses: aioresponses,
+    rest_client: Go2RtcRestClient,
+    snapshot: SnapshotAssertion,
+    additional_params: dict[str, str],
+) -> None:
+    """Test probing a stream."""
+    resp = await _test_probe(
+        responses, rest_client, "probe_success.json", 200, additional_params
+    )
+    assert resp == snapshot(name="deserialized")
+    assert isinstance(resp, Stream)
+    assert resp.to_json() == snapshot(name="serialized")
+
+
+@pytest.mark.parametrize(
+    "additional_params",
+    [
+        {"audio": "all", "video": "all"},
+        {"audio": "all"},
+        {"video": "all"},
+    ],
+    ids=[
+        "audio and video",
+        "audio only",
+        "video only",
+    ],
+)
+async def test_probe_camera_offline(
+    responses: aioresponses,
+    rest_client: Go2RtcRestClient,
+    additional_params: dict[str, str],
+) -> None:
+    """Test probing a stream, where the camera is offline."""
+    with pytest.raises(Go2RtcClientError):
+        await _test_probe(
+            responses, rest_client, "probe_camera_offline.txt", 500, additional_params
+        )