diff --git a/src/labthings_fastapi/message_broker.py b/src/labthings_fastapi/message_broker.py index d7ce98c7..17eba015 100644 --- a/src/labthings_fastapi/message_broker.py +++ b/src/labthings_fastapi/message_broker.py @@ -6,6 +6,7 @@ import logging import warnings +from collections.abc import AsyncGenerator from typing import Any, Literal from weakref import WeakSet @@ -36,7 +37,7 @@ class Message: thing: str affordance: str - message_type: Literal["property", "action"] + message_type: Literal["property", "action", "stream"] payload: Any @@ -83,6 +84,46 @@ async def subscribe( streams = affordances.setdefault(affordance, WeakSet()) streams.add(stream) + async def messages_for_affordance( + self, + thing: str, + affordance: str, + max_buffer_size: int = 1, + ) -> AsyncGenerator[Message, None]: + """Yield messages from one affordance as an async generator. + + Note that there's no timeout: standard `anyio` commands may be used to do that. + + :param thing: The name of the `.Thing` being subscribed to. + :param affordance: The name of the affordance being subscribed to. + :param max_buffer_size: The buffer size of the underlying streams. + The default will buffer up to one message. + :yield: messages from the specified affordance. + """ + send, recv = anyio.create_memory_object_stream[Message]( + max_buffer_size=max_buffer_size, + ) + await self.subscribe(thing, affordance, send) + async for message in recv: + yield message + # The stream isn't explicitly closed because it will now be discarded: the + # broker doesn't hold a strong reference. + + async def next_message(self, thing: str, affordance: str) -> Message: + """Get the next message from a particular affordance. + + Note that there's no timeout: standard `anyio` commands may be used to do that. + + :param thing: The name of the `.Thing` being subscribed to. + :param affordance: The name of the affordance being subscribed to. + :return: the next message from the specified affordance. + :raises GeneratorExit: if the stream closes before a message is received. + """ + # Note: the for loop below will do at most one iteration. + async for message in self.messages_for_affordance(thing, affordance): + return message + raise GeneratorExit("The message stream closed before anything was received.") + async def unsubscribe( self, thing: str, affordance: str, stream: MemoryObjectSendStream[Message] ) -> None: @@ -129,8 +170,13 @@ async def publish(self, message: Message) -> None: subscriptions_to_remove.add(stream) except anyio.WouldBlock: msg = f"Could not pass notification to {stream} as it was full." - LOGGER.warning(msg) - warnings.warn(MessageDroppedWarning(msg), stacklevel=1) + if message.message_type == "stream": + # It's expected that streams might drop messages, so we + # only log this at debug level. + LOGGER.debug(msg) + else: + LOGGER.warning(msg) + warnings.warn(MessageDroppedWarning(msg), stacklevel=1) for stream in subscriptions_to_remove: # discard rather than remove, so that if the stream has been finalised # since it was closed, we don't get an error. @@ -148,3 +194,22 @@ async def close_streams(self) -> None: for subs in thing_subs.values(): for stream in subs: tg.start_soon(stream.aclose) + + async def close_streams_for_affordance(self, thing: str, affordance: str) -> int: + """Close all streams subscribed to a particular affordance. + + This may be called to signal that a stream has stopped. + + :param thing: the Thing name. + :param affordance: the affordance name. + :return: the number of streams that were closed. + """ + try: + streams = self._subscriptions[thing][affordance] + except KeyError: + return 0 # If there are no streams, ignore it. + async with anyio.create_task_group() as tg: + n = len(streams) + for stream in streams: + tg.start_soon(stream.aclose) + return n diff --git a/src/labthings_fastapi/outputs/mjpeg_stream.py b/src/labthings_fastapi/outputs/mjpeg_stream.py index 39f4d7e2..3a0f64b1 100644 --- a/src/labthings_fastapi/outputs/mjpeg_stream.py +++ b/src/labthings_fastapi/outputs/mjpeg_stream.py @@ -6,33 +6,26 @@ from __future__ import annotations -import logging import threading from dataclasses import dataclass from datetime import datetime from typing import ( - TYPE_CHECKING, Any, AsyncGenerator, - Literal, Optional, - Union, - overload, ) -import anyio from fastapi import FastAPI from fastapi.responses import HTMLResponse, StreamingResponse -from typing_extensions import Self -if TYPE_CHECKING: - from labthings_fastapi.thing import Thing - from labthings_fastapi.thing_server_interface import ThingServerInterface +from labthings_fastapi.base_descriptor import BaseDescriptor +from labthings_fastapi.message_broker import Message, MessageBroker +from labthings_fastapi.thing import Thing @dataclass -class RingbufferEntry: - """A single entry in a ringbuffer. +class Frame: + """A single frame in the stream. This structure comprises one frame as a JPEG, plus a timestamp and a buffer index. Each time a frame is added to the stream, it is @@ -42,6 +35,7 @@ class RingbufferEntry: frame: bytes """The frame as a `bytes` object, which is a JPEG image for an MJPEG stream.""" + timestamp: datetime """The time the frame was captured. If the capture time is unknown then the timestamp is @@ -50,6 +44,30 @@ class RingbufferEntry: index: int """The index of the frame within the stream.""" + def __repr__(self) -> str: + """Represent the object as a string. + + :return: a representation of the object as a string. + """ + return ( + f"Frame(index={self.index}, timestamp={self.timestamp}, " + f"frame=<{len(self.frame)} bytes>)" + ) + + +def check_jpeg_bytes(frame: bytes) -> bool: + """Check a JPEG file starts and ends with the right magic bytes. + + :param frame: the data, which should be in JPEG format. + :return: `True` if the data is a valid JPEG (i.e. has valid start/end bytes). + """ + return ( + frame[0] == 0xFF + and frame[1] == 0xD8 + and frame[-2] == 0xFF + and frame[-1] == 0xD9 + ) + class MJPEGStreamResponse(StreamingResponse): """A StreamingResponse that streams an MJPEG stream. @@ -84,7 +102,7 @@ def __init__( :param status_code: The status code associated with the response, by default a 200 code is returned. """ - self.frame_async_generator = gen + self._frame_async_generator = gen StreamingResponse.__init__( self, self.mjpeg_async_generator(), @@ -101,7 +119,7 @@ async def mjpeg_async_generator(self) -> AsyncGenerator[bytes, None]: :yield: JPEG frames, each with a ``--frame`` marker prepended. """ - async for frame in self.frame_async_generator: + async for frame in self._frame_async_generator: yield b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" yield frame yield b"\r\n" @@ -128,111 +146,100 @@ class MJPEGStream: of new frames, and then retrieving the frame (shortly) afterwards. """ - def __init__( - self, thing_server_interface: ThingServerInterface, ringbuffer_size: int = 10 - ) -> None: + def __init__(self, thing: Thing, name: str) -> None: """Initialise an MJPEG stream. See the class docstring for `.MJPEGStream`. Note that it will often be initialised by `.MJPEGStreamDescriptor`. - :param thing_server_interface: the `~lt.ThingServerInterface` of the - `~lt.Thing` associated with this stream. It's used to run the async - code that relays frames to open connections. - :param ringbuffer_size: The number of frames to retain in - memory, to allow retrieval after the frame has been sent. + :param thing: the `~lt.Thing` on which this stream is defined. + :param name: The attribute name of this stream. """ self._lock = threading.Lock() - self.condition = anyio.Condition() - self._streaming = False - self._ringbuffer: list[RingbufferEntry] = [] - self._thing_server_interface = thing_server_interface - self.reset(ringbuffer_size=ringbuffer_size) + self._thing = thing + self._name = name + self.reset() - def reset(self, ringbuffer_size: Optional[int] = None) -> None: - """Reset the stream and optionally change the ringbuffer size. + @property + def _broker(self) -> MessageBroker: + """The message broker associated with our host Thing.""" + return self._thing._thing_server_interface.message_broker - Discard all frames from the ringbuffer and reset the frame index. - - :param ringbuffer_size: the number of frames to keep in memory. - """ + def reset(self) -> None: + """Reset the frame index.""" with self._lock: - self._streaming = True - n = ringbuffer_size or len(self._ringbuffer) - self._ringbuffer = [ - RingbufferEntry( - frame=b"", - index=-1, - timestamp=datetime.min, - ) - for i in range(n) - ] self.last_frame_i = -1 def stop(self) -> None: """Stop the stream. Stop the stream and cause all clients to disconnect. + This uses the thing server interface to allow it to be called from + a thread, for convenience. + + To close all streams from the event loop, use the message broker directly. """ - with self._lock: - self._streaming = False - self._thing_server_interface.start_async_task_soon( - self.notify_stream_stopped - ) + self._thing._thing_server_interface.call_async_task( + self._broker.close_streams_for_affordance, + self._thing.name, + self._name, + ) - async def ringbuffer_entry(self, i: int) -> RingbufferEntry: - """Return the ith frame acquired by the camera. + async def frames(self) -> AsyncGenerator[Frame, None]: + """Yield Frame objects (data and metadata). - The ringbuffer means we can retrieve frames even if they are not - the latest frame. Specifying ``i`` also makes it simple to ensure - that every frame in a stream is acquired. + This generator will return frames from the MJPEG stream along with timestamps, + using the `Frame` dataclass. - :param i: The index of the frame to read. + Note that this will wait for a new frame each time. There is no + guarantee that we won't skip frames. There's no timeout, if the MJPEG stream + doesn't get new frames this may block indefinitely. Frames are never buffered. + + :yield: the frames in sequence. + """ + async for message in self._broker.messages_for_affordance( + self._thing.name, + self._name, + max_buffer_size=0, + ): + if isinstance(message.payload, Frame): + # We simply ignore messages that don't contain Frame objects. + # We don't expect to get any payloads that are not Frame objects. + yield message.payload + + async def frame_async_generator(self) -> AsyncGenerator[bytes, None]: + """Yield frames as bytes objects. + + This generator will return frames from the MJPEG stream. - :return: the frame, together with a timestamp and its index. + See `frames` for how this works. This function wraps that one, returning + only the `bytes` object rather than the full `Frame` dataclass which also + includes metadata. - :raise ValueError: if the frame is not available. + :yield: the frames in sequence, as a `bytes` object containing + JPEG data. """ - if i < 0: - raise ValueError("i must be >= 0") - if i < self.last_frame_i - len(self._ringbuffer) + 2: - raise ValueError("the ith frame has been overwritten") - if i > self.last_frame_i: - # TODO: await the ith frame - raise ValueError("the ith frame has not yet been acquired") - entry = self._ringbuffer[i % len(self._ringbuffer)] - if entry.index != i: - raise ValueError("the ith frame has been overwritten") - return entry - - async def next_frame(self) -> int: - """Wait for the next frame, and return its index. - - This async function will yield until a new frame arrives, then return - its index. The index may then be used to retrieve the new frame - with `.MJPEGStream.ringbuffer_entry`. - - :return: the index of the next frame to arrive. - - :raise StopAsyncIteration: if the stream has stopped. + async for frame in self.frames(): + yield frame.frame + + async def grab_frame_with_metadata(self) -> Frame: + """Wait for the next frame, and return it. + + This includes metadata such as the timestamp. + + This method calls `frames` and returns the first frame that's yielded. + + :return: The next JPEG frame, in a dataclass with metadata. """ - async with self.condition: - await self.condition.wait() - if not self._streaming: - raise StopAsyncIteration() - return self.last_frame_i + return await anext(self.frames()) async def grab_frame(self) -> bytes: """Wait for the next frame, and return it. - This copies the frame for safety, so there is no need to release - or return the buffer. - :return: The next JPEG frame, as a `bytes` object. """ - i = await self.next_frame() - entry = await self.ringbuffer_entry(i) - return entry.frame + payload = await self.grab_frame_with_metadata() + return payload.frame async def next_frame_size(self) -> int: """Wait for the next frame and return its size. @@ -241,34 +248,7 @@ async def next_frame_size(self) -> int: :return: The size of the next JPEG frame, in bytes. """ - i = await self.next_frame() - entry = await self.ringbuffer_entry(i) - return len(entry.frame) - - async def frame_async_generator(self) -> AsyncGenerator[bytes, None]: - """Yield frames as bytes objects. - - This generator will return frames from the MJPEG stream. - - Note that this will wait for a new frame each time. There is no - guarantee that we won't skip frames. - - :yield: the frames in sequence, as a `bytes` object containing - JPEG data. - """ - while self._streaming: - try: - i = await self.next_frame() - entry = await self.ringbuffer_entry(i) - yield entry.frame - except StopAsyncIteration: - break - except Exception as e: # noqa: BLE001 - # It's important that errors in the stream don't crash the server. - # This may be something we can remove in the future, now streams stop - # more elegantly. However, it will require careful testing.f - logging.exception(f"Error in stream: {e}, stream stopped") - return + return len(await self.grab_frame()) async def mjpeg_stream_response(self) -> MJPEGStreamResponse: """Return a StreamingResponse that streams an MJPEG stream. @@ -299,48 +279,22 @@ def add_frame(self, frame: bytes, timestamp: Optional[datetime] = None) -> None: :raise ValueError: if the supplied frame does not start with the JPEG start bytes and end with the end bytes. """ - if not ( - frame[0] == 0xFF - and frame[1] == 0xD8 - and frame[-2] == 0xFF - and frame[-1] == 0xD9 - ): + if not check_jpeg_bytes(frame): raise ValueError("Invalid JPEG") with self._lock: - entry = self._ringbuffer[(self.last_frame_i + 1) % len(self._ringbuffer)] - entry.timestamp = timestamp if timestamp is not None else datetime.now() - entry.frame = frame - entry.index = self.last_frame_i + 1 - self._thing_server_interface.start_async_task_soon( - self.notify_new_frame, entry.index - ) - - async def notify_new_frame(self, i: int) -> None: - """Notify any waiting tasks that a new frame is available. - - :param i: The number of the frame (which counts up since the server starts) - """ - async with self.condition: - self.last_frame_i = i - self.condition.notify_all() - - async def notify_stream_stopped(self) -> None: - """Raise an exception in any waiting tasks to signal the stream has stopped. - - This should be run only when streaming has stopped, i.e. ``self._streaming`` - is ``False`` and an error will be raised if this isn't the case. - - :raises RuntimeError: if the stream is still streaming. - """ - if self._streaming is True: - raise RuntimeError( - "This function should only be called when the stream is stopped." - ) - async with self.condition: - self.condition.notify_all() + self.last_frame_i += 1 + frame_index = self.last_frame_i + payload = Frame( + frame=frame, + timestamp=timestamp if timestamp is not None else datetime.now(), + index=frame_index, + ) + self._thing._thing_server_interface.publish( + Message(self._thing.name, self._name, "stream", payload) + ) -class MJPEGStreamDescriptor: +class MJPEGStreamDescriptor(BaseDescriptor[Thing, MJPEGStream]): """A descriptor that returns a MJPEGStream object when accessed. If this descriptor is added to a `~lt.Thing`, it will create an `.MJPEGStream` @@ -357,48 +311,23 @@ def __init__(self, **kwargs: Any) -> None: :param \**kwargs: keyword arguments are passed to the initialiser of `.MJPEGStream`. """ + super().__init__() self._kwargs: Any = kwargs - def __set_name__(self, _owner: Thing, name: str) -> None: - """Remember the name to which we are assigned. - - The name is important, as it will set the URL of the HTTP endpoint used - to access the stream. - - :param _owner: the `~lt.Thing` to which we are attached. - :param name: the name to which this descriptor is assigned. - """ - self.name = name - - @overload - def __get__(self, obj: Literal[None], type: type | None = None) -> Self: ... # noqa: D105 - - @overload - def __get__(self, obj: Thing, type: type | None = None) -> MJPEGStream: ... # noqa: D105 - - def __get__( - self, obj: Optional[Thing], type: type[Thing] | None = None - ) -> Union[MJPEGStream, Self]: - """Return the MJPEG Stream, or the descriptor object. - - When accessed on the class, this ``__get__`` method will return the descriptor - object. This allows LabThings to add it to the HTTP API. - - When accessed on the object, an `.MJPEGStream` is returned. + def instance_get(self, obj: Thing) -> MJPEGStream: + """Return the MJPEG Stream. - :param obj: the host `~lt.Thing`, or ``None`` if accessed on the class. - :param type: the class on which we are defined. + :param obj: the host `~lt.Thing`. :return: an `.MJPEGStream`, or this descriptor. """ - if obj is None: - return self try: return obj.__dict__[self.name] except KeyError: obj.__dict__[self.name] = MJPEGStream( + thing=obj, + name=self.name, **self._kwargs, - thing_server_interface=obj._thing_server_interface, ) return obj.__dict__[self.name] diff --git a/src/labthings_fastapi/thing_server_interface.py b/src/labthings_fastapi/thing_server_interface.py index 3acbf011..a9d0bb8c 100644 --- a/src/labthings_fastapi/thing_server_interface.py +++ b/src/labthings_fastapi/thing_server_interface.py @@ -20,7 +20,7 @@ from labthings_fastapi.exceptions import FeatureNotEnabledError, ServerNotRunningError from labthings_fastapi.global_lock import GlobalLock -from labthings_fastapi.message_broker import Message +from labthings_fastapi.message_broker import Message, MessageBroker if TYPE_CHECKING: from labthings_fastapi.actions import ActionManager @@ -156,6 +156,11 @@ def publish(self, message: Message) -> None: except ServerNotRunningError: pass # If the server isn't running yet, we can't publish events. + @property + def message_broker(self) -> MessageBroker: + """The message broker associated with the current server.""" + return self._get_server().message_broker + @property def settings_folder(self) -> str: """The path to a folder where persistent files may be saved.""" diff --git a/tests/test_message_broker.py b/tests/test_message_broker.py index dbfea33d..e3de8925 100644 --- a/tests/test_message_broker.py +++ b/tests/test_message_broker.py @@ -4,6 +4,7 @@ from weakref import WeakSet import anyio +import anyio.lowlevel import pytest from anyio.abc import ObjectReceiveStream from pydantic import ValidationError @@ -23,6 +24,7 @@ class Unjsonable: ("test_thing", "prop", "property", Unjsonable()), ("test_thing", "do_it", "action", None), ("test_thing", "notify", "action", {"key": "value"}), + ("test_thing", "mjpeg", "stream", b"whatever"), ], ) def test_message_valid(message): @@ -91,6 +93,23 @@ async def test_subscribe_unsubscribe(): await broker.unsubscribe("whatever", 42, send_stream) # type: ignore +async def test_unsubscribe_on_deletion(): + """Check that streams auto-unsubscribe when they are finalised.""" + broker = MessageBroker() + assert broker._subscriptions == {} + + # subscribe to a stream + send_stream, receive_stream = anyio.create_memory_object_stream[Message]() + await broker.subscribe("thing_name", "prop", send_stream) + assert send_stream in broker._subscriptions["thing_name"]["prop"] + + # delete the streams + del send_stream + del receive_stream + + assert len(broker._subscriptions["thing_name"]["prop"]) == 0 + + async def append_messages( stream: ObjectReceiveStream[Message], dest: list[Message], @@ -186,7 +205,6 @@ async def test_sending_to_closed_streams(caplog, action): assert send_stream in broker._subscriptions["thing"]["prop"] message = Message("thing", "prop", "property", None) - print("sending first message") # Check we can send and receive a message await broker.publish(message) received = await receive_stream.receive() @@ -200,7 +218,6 @@ async def test_sending_to_closed_streams(caplog, action): else: del send_stream # streams are unsubscribed when they are finalised assert len(caplog.records) == 0 - print("sending second message") await broker.publish(message) assert len(caplog.records) == 0 # Shouldn't be any warnings in the log # Check we've been unsubscribed @@ -228,6 +245,19 @@ async def test_sending_to_full_stream(caplog): with pytest.warns(MessageDroppedWarning): await broker.publish(message) assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.WARNING + msg = caplog.records[0].getMessage() + assert msg.startswith("Could not pass notification to") + assert msg.endswith("as it was full.") + caplog.clear() + + # Send a fourth message marked as being part of a stream. + # This should only log at DEBUG level. + stream_msg = Message("thing", "prop", "stream", None) + caplog.set_level(logging.DEBUG) + await broker.publish(stream_msg) + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.DEBUG msg = caplog.records[0].getMessage() assert msg.startswith("Could not pass notification to") assert msg.endswith("as it was full.") @@ -260,3 +290,100 @@ async def test_weakset_garbage_collection(): assert len(iterated) == 1 # We only iterated once, but there wasn't an error. # If we complete the test, it confirms we don't need to worry about streams being # finalised during iteration. + + +async def collect_messages(generator, limit: int | None = None) -> list[Message]: + """Collect messages from a generator into a list.""" + received = [] + async for message in generator: + received.append(message) + if limit and len(received) >= limit: + break + return received + + +@pytest.mark.parametrize("buffer_size", (0, 1, 3)) +@pytest.mark.parametrize("close_stream", ("subscriber", "broker")) +async def test_messages_for_affordance(buffer_size, close_stream): + """Test the shorthand to iterate over messages for an affordance works. + + This also tests close_streams_for_affordance(). + + Note that there are various uses of `await anyio.lowlevel.checkpoint()` in + this test. It's necessary to add these in order to yield to the event + loop so that the task that listens is able to run. + + For buffer sizes from 0 to 3, we send (b+1) messages at which point the + buffer is full. We then send another message that is dropped. + + After sending the messages, we then await (b+2) checkpoints, to allow + the task to clear the buffer, and try again. + """ + broker = MessageBroker() + messages = [Message("thing", "prop", "property", i) for i in range(buffer_size + 1)] + dropped_message = Message("thing", "prop", "property", "dropped") + + # We subscribe to the broker in a background task, and send messages. + # With no buffer, we should process one message each time we yield to the + # event loop. With a finite buffer, we process (b+1) messages each time. + # We make b+2 messages, and check we get the first (b+1) of them. + async with anyio.create_task_group() as tg: + # Start a background task to listen for messages, using a pair of streams + # with the specified buffer size. + handle = tg.start_soon( + collect_messages, + broker.messages_for_affordance("thing", "prop", buffer_size), + 2 * (buffer_size + 1) if close_stream == "subscriber" else None, + ) + # Allow the background task to run (subscribe and start listening) + await anyio.lowlevel.checkpoint() + await anyio.lowlevel.checkpoint() + # We loop around twice, sending all the messages + for _ in range(2): + for message in messages: + await broker.publish(message) + # After sending these, the buffer should be full + with pytest.warns(MessageDroppedWarning): + await broker.publish(dropped_message) + # Yield to the event loop, so the task can clear its buffer. + for _ in range(buffer_size + 2): + await anyio.lowlevel.checkpoint() + if close_stream == "broker": + # The background task will run until its stream is closed. + # the next line should close the stream: if the test doesn't + # hang, it's worked. + await broker.close_streams_for_affordance("thing", "prop") + # Calling this twice shouldn't error. + await broker.close_streams_for_affordance("thing", "prop") + # Each time, we should have dropped the last message (because we send two + # more than the buffer size). + assert handle.return_value == messages * 2 + # As the listener stopped, its stream should no longer be subscribed + assert len(broker._subscriptions["thing"]["prop"]) == 0 + + +async def test_next_message(): + """Test the shorthand to get the next message for an affordance. + + This also tests close_streams_for_affordance(). + """ + broker = MessageBroker() + message = Message("thing", "prop", "property", None) + + # Listen for three messages (and then stop listening) + async with anyio.create_task_group() as tg: + handle = tg.start_soon(broker.next_message, "thing", "prop") + await anyio.lowlevel.checkpoint() + await broker.publish(message) + assert handle.return_value == message + # The stream should no longer be subscribed + assert len(broker._subscriptions["thing"]["prop"]) == 0 + + # Test for error if the stream closes without a message + with pytest.RaisesGroup(GeneratorExit, flatten_subgroups=True): + async with anyio.create_task_group() as tg: + handle = tg.start_soon(broker.next_message, "thing", "prop") + await anyio.lowlevel.checkpoint() + await broker.close_streams_for_affordance("thing", "prop") + # The stream should no longer be subscribed + assert len(broker._subscriptions["thing"]["prop"]) == 0 diff --git a/tests/test_mjpeg_stream.py b/tests/test_mjpeg_stream.py index 1aefa2b8..1fc93039 100644 --- a/tests/test_mjpeg_stream.py +++ b/tests/test_mjpeg_stream.py @@ -1,18 +1,38 @@ import io +import logging import threading import time +from datetime import datetime +import anyio import pytest +from anyio import from_thread, to_thread from PIL import Image import labthings_fastapi as lt +from labthings_fastapi.message_broker import Message, MessageBroker +from labthings_fastapi.outputs.mjpeg_stream import ( + Frame, + MJPEGStream, + check_jpeg_bytes, +) + + +def make_jpeg(colour: str) -> bytes: + """Return a 10x10 JPEG of the given colour.""" + image = Image.new("RGB", (10, 10), colour) + dest = io.BytesIO() + image.save(dest, "jpeg") + return dest.getvalue() class Telly(lt.Thing): _stream_thread: threading.Thread _streaming: bool = False framerate: float = 1000 - frame_limit: int = 3 + frame_limit: int = 999 + frame_event: threading.Event | None = None + initial_delay: float = 0 stream = lt.outputs.MJPEGStreamDescriptor() @@ -23,60 +43,333 @@ def __enter__(self): def __exit__(self, exc_t, exc_v, exc_tb): self._streaming = False + if self.frame_event: + # Trigger an iteration of the loop, so that it + # will terminate rather than hang forever + self.frame_event.set() 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()) + jpegs = [make_jpeg(c) for c in colours] + + if self.initial_delay > 0: + time.sleep(self.initial_delay) i = 0 while self._streaming and (i < self.frame_limit or self.frame_limit < 0): self.stream.add_frame(jpegs[i % len(jpegs)]) - time.sleep(1 / self.framerate) i = i + 1 + if self.frame_event: + self.frame_event.wait() + self.frame_event.clear() + else: + time.sleep(1 / self.framerate) self.stream.stop() self._streaming = False +def assert_magic_bytes(frame: bytes) -> None: + """Check that a `bytes` object starts and ends with the JPEG markers.""" + assert frame[0] == 0xFF + assert frame[1] == 0xD8 + assert frame[-2] == 0xFF + assert frame[-1] == 0xD9 + + @pytest.fixture -def client(): - """Yield a test client connected to a ThingServer""" - server = lt.ThingServer.from_things({"telly": Telly}) - with server.test_client() as client: - yield client +def server(): + """Yield a ThingServer with the `Telly` thing.""" + return lt.ThingServer.from_things({"telly": Telly}) + + +@pytest.fixture +def telly(server): + """Yield a Telly thing from the server.""" + telly = server.things["telly"] + assert isinstance(telly, Telly) + return telly + + +def test_check_jpeg_bytes(jpeg_bytes): + """Check the check_jpeg_bytes function works correctly""" + assert check_jpeg_bytes(jpeg_bytes) is True + assert check_jpeg_bytes(b"00000000") is False + + +def test_frame_repr(jpeg_bytes): + """Make sure Frame objects don't print JPEG data.""" + frame = Frame(jpeg_bytes, datetime(2005, 7, 14, 12, 30), 1) + # I'll be explicit here: I want to make sure there's no dump + # of the bytes. + assert 'b"' not in repr(frame) + # I can also just test the known repr: + assert repr(frame) == ( + "Frame(index=1, timestamp=2005-07-14 12:30:00, " + f"frame=<{len(jpeg_bytes)} bytes>)" + ) + assert str(frame) == repr(frame) + + +def test_grab_and_shutdown(server: lt.ThingServer, telly: Telly): + """Check we can grab frames, and shut down cleanly. + + This test uses an Event to synchronise new frames with the various + methods we're calling to retrieve them. This is intended to make the + test suite faster and less reliant on `time.sleep`. + + We check that `grab_frame` and `next_frame_size` both work when the + camera emits frames, and also that they raise an error if they're called + once the camera has stopped. + + We verify that the stream can be stopped, though we don't shut down any + long-running listeners, because limitations in TestClient mean this + isn't possible. It would be possible if we spun up an actual HTTP server, + but that's quite high-effort. + + The async grab functions need to run in the event loop, which happens in + a background thread during the `with server.test_client()` block. We use + the thing server interface to run tasks in the event loop for convenience. + """ + telly.frame_event = threading.Event() # Make timings more deterministic. + with server.test_client(): + # this `with` block starts an event loop and runs the server. + + # Grab some frames and check we get a JPEG back + for _ in range(3): + assert telly._stream_thread.is_alive() # Catch premature termination + # Start the coroutine to grab a frame + future = telly._thing_server_interface.start_async_task_soon( + telly.stream.grab_frame + ) + # then use the event to trigger the next frame. + telly.frame_event.set() + # then wait for the coroutine to finish + frame = future.result() + # this should return a valid JPEG, which we can (sort of) check below + assert_magic_bytes(frame) + + # Repeat the process for grabbing a frame to test `next_frame_size` + future = telly._thing_server_interface.start_async_task_soon( + telly.stream.next_frame_size + ) + telly.frame_event.set() + size = future.result() # wait for the frame to be grabbed + assert isinstance(size, int) + assert size > 0 + + # Close all streams + telly.stream.stop() + + # The background thread gets shut down by `Telly.__exit__`. + + +@pytest.fixture +def broker() -> MessageBroker: + """A MessageBroker object.""" + return MessageBroker() + + +@pytest.fixture +def mjpeg_stream(mocker, broker): + """An MJPEGStream object with a mocked Thing.""" + thing = mocker.Mock() + tsi = thing._thing_server_interface + tsi.message_broker = broker + tsi.start_async_task_soon.side_effect = from_thread.run + + def publish(message: Message) -> None: + from_thread.run(broker.publish, message) + tsi.publish.side_effect = publish + thing.name = "thing" + yield MJPEGStream(thing, "stream") -def test_mjpeg_stream(client): - """Verify the MJPEG stream contains at least one frame marker. + +@pytest.fixture +def jpeg_bytes(): + return make_jpeg("#F0F") + + +async def test_grab_frame(mjpeg_stream, jpeg_bytes): + """Test we can add a frame in a thread and grab its bytes.""" + # Listen for the next frame (grab_frame) then add a frame. + async with anyio.create_task_group() as tg: + handle = tg.start_soon(mjpeg_stream.grab_frame) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + assert handle.return_value is jpeg_bytes + + +async def test_grab_frame_with_metadata(mjpeg_stream, jpeg_bytes): + """Test we can add a frame in a thread and grab it.""" + # Listen for the next frame (grab_frame) then add a frame. + async with anyio.create_task_group() as tg: + handle = tg.start_soon(mjpeg_stream.grab_frame_with_metadata) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + assert handle.return_value.frame is jpeg_bytes + + +async def test_next_frame_size(mjpeg_stream, jpeg_bytes): + """Test we can add a frame in a thread and grab its size.""" + # Listen for the next frame (grab_frame) then add a frame. + async with anyio.create_task_group() as tg: + handle = tg.start_soon(mjpeg_stream.next_frame_size) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + assert handle.return_value == len(jpeg_bytes) + + +async def grab_three_frames_with_generator(stream: MJPEGStream) -> list[Frame]: + """Grab three frames and return them in a list.""" + frames = [] + async for frame in stream.frames(): + frames.append(frame) + if len(frames) >= 3: + break + return frames + + +async def test_frames(mjpeg_stream, jpeg_bytes, broker): + """Check we can get a few frames from the generator. + + This test also makes sure we don't end up with dangling + streams that are no longer wanted. + """ + async with anyio.create_task_group() as tg: + handle = tg.start_soon(grab_three_frames_with_generator, mjpeg_stream) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + frames = handle.return_value + assert len(frames) == 3 + for i, f in enumerate(frames): + assert isinstance(f, Frame) + assert f.frame is jpeg_bytes + assert f.index == i + # Emit another frame, to ensure it deletes the stream + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + # Check there's no streams left connected + assert len(broker._subscriptions["thing"]["stream"]) == 0 + + +async def grab_three_frames_from_old_generator(stream: MJPEGStream) -> list[bytes]: + """Grab three frames from the old-style generator""" + frames = [] + async for frame in stream.frame_async_generator(): + frames.append(frame) + if len(frames) >= 3: + break + return frames + + +async def test_frame_async_generator(mjpeg_stream, jpeg_bytes, broker): + """Check we can get a few frames from the old generator. + + This test also makes sure we don't end up with dangling + streams that are no longer wanted. + """ + async with anyio.create_task_group() as tg: + handle = tg.start_soon(grab_three_frames_from_old_generator, mjpeg_stream) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + frames = handle.return_value + assert len(frames) == 3 + for f in frames: + assert isinstance(f, bytes) + assert f is jpeg_bytes + # Emit another frame, to ensure it deletes the stream + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + # Check there's no streams left connected + assert len(broker._subscriptions["thing"]["stream"]) == 0 + + +async def test_dropped_frames(mjpeg_stream, jpeg_bytes, broker, caplog): + """Test what happens when a subscribed stream is full. + + Note that the MessageBroker functionality is also tested in + test_message_broker.py::test_sending_to_full_stream. + """ + send, recv = anyio.create_memory_object_stream[Message](1) + await broker.subscribe("thing", "stream", send) + caplog.set_level(logging.DEBUG) + caplog.clear() + # The first message sits in the buffer, the second gets dropped + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + await to_thread.run_sync(mjpeg_stream.add_frame, jpeg_bytes) + print(caplog.records) + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.DEBUG + msg = caplog.records[0].getMessage() + assert msg.startswith("Could not pass notification to") + assert msg.endswith("as it was full.") + # We should summarise the length of the JPEG, not include + # the JPEG bytes. + assert f"<{len(jpeg_bytes)} bytes>" in msg + # We want to make sure we're not dumping JPEGs to the log + assert 'b"' not in msg + + +def test_mjpeg_stream_http(server: lt.ThingServer, telly: Telly): + """Verify the MJPEG stream works, and is shut down cleanly. A limitation of the TestClient is that it can't actually stream. This means that all of the frames sent by our test Thing will arrive in a single packet. - For now, we just check it starts with the frame separator, - but it might be possible in the future to check there are three - images there. + For now, we download all the data and then chop it up afterwards. + + The `Telly` will send exactly 3 JPEGs, with a short delay at + the start to make sure the `StreamingResponse` is created and + doesn't miss any frames. + + This test also verifies the stream is shut down by the server + when it's stopped - if it wasn't terminated by the server, the + `client.stream()` call would hang indefinitely. """ - with client.stream("GET", "/telly/stream") as stream: - stream.raise_for_status() - received = 0 - for b in stream.iter_bytes(): - received += 1 - assert b.startswith(b"--frame") + telly.frame_limit = 3 # stream 3 frames and then stop + telly.initial_delay = 0.05 # give the client time to start listening + with server.test_client() as client: + with client.stream("GET", "/telly/stream") as response: + # Note: we don't actually enter this `with` block until after + # the camera background thread has finished and the stream + # has been closed. + response.raise_for_status() + parts = 0 + mjpeg_data = b"" + for b in response.iter_bytes(): + parts += 1 + mjpeg_data += b + # Due to a quirk in TestClient, we get all the data in a single + # chunk - this limits our ability to test the stream. + # If that's fixed in the future, the assertion below will fail, + # which should prompt us to improve these tests. + assert parts == 1 + # Check the received data contained the expected number of frames + # We split chunks using the known header, and remove extra white + # space before checking for JPEG start/end bytes. + chunks = mjpeg_data.split(b"--frame\r\nContent-Type: image/jpeg") + n = 0 + for chunk in chunks: + # Check each chunk is a JPEG + stripped = chunk.strip(b"\r\n") + if len(stripped) > 10: # If the chunk doesn't look empty + assert_magic_bytes(stripped) + n += 1 + assert telly.stream.last_frame_i == 2 + assert n == telly.frame_limit -if __name__ == "__main__": - import uvicorn - server = lt.ThingServer.from_things({"telly": Telly}) - telly = server.things["telly"] +if __name__ == "__main__": + """This block allows you to connect manually with a web browser. + + That's helpful, because the tests above don't actually stream anything + as noted in `test_mjpeg_stream_http`. + """ + thing_server = lt.ThingServer.from_things({"telly": Telly}) + telly = thing_server.things["telly"] assert isinstance(telly, Telly) telly.framerate = 6 telly.frame_limit = -1 - uvicorn.run(server.app, port=5000, ws="websockets-sansio") + thing_server.serve()