WIP: Add a test for the MJPEG Stream

Currently this fails, because the `TestClient` doesn't actually stream responses :(

A work-around might be to implement stopping (i.e. get it to raise an exception and close the stream), then check the response includes some frames.

Author: rwb27

pyproject.toml

@@ -30,6 +30,7 @@ dev = [
   "mypy>=1.6.1, <2",
   "ruff>=0.1.3",
   "types-jsonschema",
+  "Pillow",
 ]
 
 [project.urls]

src/labthings_fastapi/outputs/mjpeg_stream.py

@@ -41,7 +41,11 @@ def __init__(self, gen: AsyncGenerator[bytes, None], status_code: int = 200):
 
         This response is initialised with an async generator that yields `bytes`
         objects, each of which is a JPEG file. We add the --frame markers and mime
-        types that enable it to work in an `img` tag.
+        types that mark it as an MJPEG stream. This is sufficient to enable it to
+        work in an `img` tag, with the `src` set to the MJPEG stream's endpoint.
+
+        It expects an async generator that supplies individual JPEGs to be streamed,
+        such as the one provided by `.MJPEGStream`.
 
         NB the ``status_code`` argument is used by FastAPI to set the status code of
         the response in OpenAPI.
@@ -63,6 +67,24 @@ async def mjpeg_async_generator(self) -> AsyncGenerator[bytes, None]:
 
 
 class MJPEGStream:
+    """Manage streaming images over HTTP as an MJPEG stream
+
+    An MJPEGStream object handles accepting images (already in
+    JPEG format) and streaming them to HTTP clients as a multipart
+    response.
+
+    The minimum needed to make the stream work is to periodically
+    call `add_frame` with JPEG image data.
+
+    To add a stream to a `.Thing`, use the `.MJPEGStreamDescriptor`
+    which will handle creating an `MJPEGStream` object on first access,
+    and will also add it to the HTTP API.
+
+    The MJPEG stream buffers the last few frames (10 by default) and
+    also has a hook to notify the size of each frame as it is added.
+    The latter is used by OpenFlexure's autofocus routine.
+    """
+
     def __init__(self, ringbuffer_size: int = 10):
         self._lock = threading.Lock()
         self.condition = anyio.Condition()
@@ -182,7 +204,14 @@ async def notify_new_frame(self, i):
 
 
 class MJPEGStreamDescriptor:
-    """A descriptor that returns a MJPEGStream object when accessed"""
+    """A descriptor that returns a MJPEGStream object when accessed
+
+    If this descriptor is added to a `.Thing`, it will create an `.MJPEGStream`
+    object when it is first accessed. It will also add two HTTP endpoints,
+    one with the name of the descriptor serving the MJPEG stream, and another
+    with `/viewer` appended, which serves a basic HTML page that views the stream.
+
+    """
 
     def __init__(self, **kwargs):
         self._kwargs = kwargs

tests/test_mjpeg_stream.py

@@ -0,0 +1,66 @@
+import io
+import threading
+import time
+from PIL import Image
+from fastapi.testclient import TestClient
+import labthings_fastapi as lt
+
+
+class Telly(lt.Thing):
+    _stream_thread: threading.Thread
+    _streaming: bool = False
+    framerate: float = 1000
+
+    stream = lt.outputs.MJPEGStreamDescriptor()
+
+    def __enter__(self):
+        self._streaming = True
+        self._stream_thread = threading.Thread(target=self._make_images)
+        self._stream_thread.start()
+
+    def __exit__(self, exc_t, exc_v, exc_tb):
+        self._streaming = False
+        self._stream_thread.join()
+
+    def _make_images(self):
+        """Stream a series of solid colours"""
+        colours = ["#F00", "#0F0", "#00F"]
+        jpegs = []
+        for c in colours:
+            image = Image.new("RGB", (10, 10), c)
+            dest = io.BytesIO()
+            image.save(dest, "jpeg")
+            jpegs.append(dest.getvalue())
+
+        i = -1
+        start_time = time.time()
+        while self._streaming:
+            i = (i + 1) % len(jpegs)
+            print(f"sending frame {i}")
+            self.stream.add_frame(jpegs[i], self._labthings_blocking_portal)
+            time.sleep(1 / self.framerate)
+
+            if time.time() - start_time > 10:
+                break
+        print("stopped sending frames")
+        self._streaming = False
+
+
+def test_mjpeg_stream():
+    server = lt.ThingServer()
+    telly = Telly()
+    server.add_thing(telly, "telly")
+    with TestClient(server.app) as client:
+        with client.stream("GET", "/telly/stream", timeout=0.1) as stream:
+            stream.raise_for_status()
+            received = 0
+            for b in stream.iter_bytes():
+                received += 1
+                print(f"Got packet {received}")
+                assert b.startswith(b"--frame")
+                if received > 5:
+                    break
+
+
+if __name__ == "__main__":
+    test_mjpeg_stream()