From aba153eecd28bbeb383a2e917fa65e109572e0a1 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Thu, 9 Jul 2026 09:45:20 +0300 Subject: [PATCH] ustreamer: video streaming j video Usage: j video [OPTIONS] COMMAND [ARGS]... Video capture and streaming Options: --help Show this message and exit. Commands: snapshot Save a single snapshot to file state Show video source state stream Start local MJPEG streaming server Signed-off-by: Benny Zlotnik Assisted-by: grok-4.5 --- .../jumpstarter-driver-ustreamer/README.md | 65 ++++ .../jumpstarter_driver_ustreamer/client.py | 162 ++++++++- .../client_test.py | 318 ++++++++++++++++++ .../jumpstarter_driver_ustreamer/driver.py | 36 +- .../driver_test.py | 62 +++- 5 files changed, 634 insertions(+), 9 deletions(-) create mode 100644 python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client_test.py diff --git a/python/packages/jumpstarter-driver-ustreamer/README.md b/python/packages/jumpstarter-driver-ustreamer/README.md index 22b0f3737..f608ec839 100644 --- a/python/packages/jumpstarter-driver-ustreamer/README.md +++ b/python/packages/jumpstarter-driver-ustreamer/README.md @@ -28,6 +28,71 @@ Traceback (most recent call last): io.UnsupportedOperation: fileno ``` +## Usage + +### CLI + +The uStreamer client exposes a `video` command inside `jmp shell` for +inspecting the current stream state, saving a snapshot, and starting a local +MJPEG proxy server. + +```console +$ j video --help +Usage: j video [OPTIONS] COMMAND [ARGS]... + + Video capture and streaming + +Options: + --help Show this message and exit. + +Commands: + snapshot Save a single snapshot to file + state Show video source state + stream Start local MJPEG streaming server +``` + +#### `j video state` + +```console +$ j video state --help +Usage: j video state [OPTIONS] + + Show video source state + +Options: + --help Show this message and exit. +``` + +#### `j video snapshot` + +```console +$ j video snapshot --help +Usage: j video snapshot [OPTIONS] + + Save a single snapshot to file + +Options: + -o, --output TEXT Output file path + --help Show this message and exit. +``` + +#### `j video stream` + +```console +$ j video stream --help +Usage: j video stream [OPTIONS] + + Start local MJPEG streaming server + + Proxies ustreamer's native MJPEG stream through the jumpstarter tunnel. Frame + rate is controlled by ustreamer's configuration. + +Options: + -p, --port INTEGER Local server port (0 = auto) + --browser / --no-browser Open in web browser + --help Show this message and exit. +``` + ## API Reference ```{eval-rst} diff --git a/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client.py b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client.py index 2be00db0a..1365cbd08 100644 --- a/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client.py +++ b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client.py @@ -1,10 +1,106 @@ import io +import webbrowser from base64 import b64decode +import click +from aiohttp import web +from anyio import EndOfStream, get_cancelled_exc_class, move_on_after from PIL import Image from .common import UStreamerState from jumpstarter.client import DriverClient +from jumpstarter.client.decorators import driver_click_group + +LANDING_PAGE = """\ + + + + Video + + + +

Jumpstarter Video Stream

+ Live video stream +

Single snapshot (JPEG)

+ + +""" + + +def _parse_content_type(header_bytes: bytes) -> str: + """Extract Content-Type from raw HTTP response headers.""" + for line in header_bytes.decode("ascii", errors="replace").split("\r\n"): + if line.lower().startswith("content-type:"): + return line.split(":", 1)[1].strip() + return "multipart/x-mixed-replace; boundary=--" + + +def _run_server(client, app, port, open_browser): + """Run an aiohttp app, opening the browser and blocking until Ctrl+C.""" + runner = web.AppRunner(app) + + async def serve(): + await runner.setup() + try: + site = web.TCPSite(runner, "127.0.0.1", port) + await site.start() + + addresses = runner.addresses + if not addresses: + raise RuntimeError("Video server started without a bound address") + actual_port = int(addresses[0][1]) + url = f"http://127.0.0.1:{actual_port}" + click.echo(f"Video stream available at: {url}") + click.echo(f"Snapshot endpoint: {url}/snapshot") + click.echo("Press Ctrl+C to stop.") + + if open_browser: + webbrowser.open(url) + + from anyio import sleep_forever + await sleep_forever() + finally: + with move_on_after(2, shield=True): + await runner.cleanup() + + try: + client.portal.call(serve) + except KeyboardInterrupt: + click.echo("\nStopping video server.") + + +async def _proxy_mjpeg_stream(client, request): + """Proxy ustreamer's native MJPEG stream through the jumpstarter tunnel.""" + async with client.stream_async("connect") as tunnel: + await tunnel.send(b"GET /stream HTTP/1.1\r\nHost: localhost\r\n\r\n") + + buf = b"" + while b"\r\n\r\n" not in buf: + buf += await tunnel.receive() + + header_part, _, body_start = buf.partition(b"\r\n\r\n") + + response = web.StreamResponse() + response.content_type = _parse_content_type(header_part) + await response.prepare(request) + + if body_start: + await response.write(body_start) + + try: + while True: + chunk = await tunnel.receive() + await response.write(chunk) + except (EndOfStream, ConnectionResetError, ConnectionAbortedError, get_cancelled_exc_class()): + pass + + return response class UStreamerClient(DriverClient): @@ -14,18 +110,72 @@ class UStreamerClient(DriverClient): """ def state(self): - """ - Get state of ustreamer service - """ - + """Get state of ustreamer service""" return UStreamerState.model_validate(self.call("state")) def snapshot(self): - """ - Get a snapshot image from the video input + """Get a snapshot image from the video input :return: PIL Image object of the snapshot image :rtype: PIL.Image """ input_jpg_data = b64decode(self.call("snapshot")) return Image.open(io.BytesIO(input_jpg_data)) + + def snapshot_bytes(self): + """Get raw JPEG bytes from the video input""" + return b64decode(self.call("snapshot")) + + def cli(self): + @driver_click_group(self) + def video(): + """Video capture and streaming""" + pass + + @video.command() + def state(): + """Show video source state""" + s = self.state() + src = s.result.source + enc = s.result.encoder + click.echo(f"Online: {src.online}") + click.echo(f"Resolution: {src.resolution.width}x{src.resolution.height}") + click.echo(f"FPS: {src.captured_fps}/{src.desired_fps}") + click.echo(f"Encoder: {enc.type} (quality: {enc.quality})") + + @video.command() + @click.option("-o", "--output", default="snapshot.jpg", help="Output file path") + def snapshot(output): + """Save a single snapshot to file""" + img = self.snapshot() + img.save(output) + click.echo(f"Saved snapshot to {output}") + + @video.command() + @click.option("-p", "--port", default=0, type=int, help="Local server port (0 = auto)") + @click.option("--browser/--no-browser", default=True, help="Open in web browser") + def stream(port, browser): + """Start local MJPEG streaming server + + Proxies ustreamer's native MJPEG stream through the jumpstarter + tunnel. Frame rate is controlled by ustreamer's configuration. + """ + + async def handle_index(request): + return web.Response(text=LANDING_PAGE, content_type="text/html") + + async def handle_snapshot(request): + data = b64decode(await self.call_async("snapshot")) + return web.Response(body=data, content_type="image/jpeg") + + async def handle_stream(request): + return await _proxy_mjpeg_stream(self, request) + + app = web.Application() + app.router.add_get("/", handle_index) + app.router.add_get("/snapshot", handle_snapshot) + app.router.add_get("/stream", handle_stream) + + _run_server(self, app, port, browser) + + return video diff --git a/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client_test.py b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client_test.py new file mode 100644 index 000000000..e5840f2d3 --- /dev/null +++ b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/client_test.py @@ -0,0 +1,318 @@ +import base64 +import contextlib +import io +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import anyio +import pytest +from click.testing import CliRunner +from PIL import Image + +from jumpstarter_driver_ustreamer.client import ( + LANDING_PAGE, + UStreamerClient, + _parse_content_type, + _proxy_mjpeg_stream, + _run_server, +) +from jumpstarter_driver_ustreamer.common import UStreamerState + + +def _make_client(): + client = object.__new__(UStreamerClient) + client.description = None + client.methods_description = {} + client.stack = MagicMock() + return client + + +def _make_state(): + return UStreamerState.model_validate( + { + "ok": True, + "result": { + "source": { + "online": True, + "desired_fps": 30, + "captured_fps": 29, + "resolution": {"width": 1280, "height": 720}, + }, + "encoder": {"type": "CPU", "quality": 85}, + }, + } + ) + + +def _make_jpeg_bytes(): + buffer = io.BytesIO() + Image.new("RGB", (2, 1), color="red").save(buffer, format="JPEG") + return buffer.getvalue() + + +def _get_route_handler(app, path): + for resource in app.router.resources(): + if getattr(resource, "canonical", None) == path: + return next(iter(resource)).handler + raise AssertionError(f"Route {path} was not registered") + + +class _FakeStreamContext: + def __init__(self, tunnel): + self.tunnel = tunnel + + async def __aenter__(self): + return self.tunnel + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class _FakeTunnel: + def __init__(self, chunks): + self._chunks = iter(chunks) + self.sent = [] + + async def send(self, data): + self.sent.append(data) + + async def receive(self): + chunk = next(self._chunks) + if isinstance(chunk, BaseException): + raise chunk + return chunk + + +class _FakeStreamResponse: + def __init__(self): + self.content_type = None + self.prepared_request = None + self.writes = [] + + async def prepare(self, request): + self.prepared_request = request + + async def write(self, data): + self.writes.append(data) + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +def test_parse_content_type_returns_header_value(): + header_bytes = b"HTTP/1.1 200 OK\r\nContent-Type: image/jpeg\r\nX-Test: 1\r\n" + + assert _parse_content_type(header_bytes) == "image/jpeg" + + +def test_parse_content_type_falls_back_to_mjpeg_default(): + assert _parse_content_type(b"HTTP/1.1 200 OK\r\nX-Test: 1\r\n") == "multipart/x-mixed-replace; boundary=--" + + +def test_state_returns_validated_model(): + client = _make_client() + client.call = MagicMock(return_value=_make_state().model_dump(mode="python")) + + state = UStreamerClient.state(client) + + assert state == _make_state() + client.call.assert_called_once_with("state") + + +def test_snapshot_returns_image(): + client = _make_client() + jpeg_bytes = _make_jpeg_bytes() + client.call = MagicMock(return_value=base64.b64encode(jpeg_bytes).decode("ascii")) + + image = UStreamerClient.snapshot(client) + + assert image.size == (2, 1) + client.call.assert_called_once_with("snapshot") + + +def test_snapshot_bytes_returns_decoded_jpeg(): + client = _make_client() + jpeg_bytes = _make_jpeg_bytes() + client.call = MagicMock(return_value=base64.b64encode(jpeg_bytes).decode("ascii")) + + assert UStreamerClient.snapshot_bytes(client) == jpeg_bytes + client.call.assert_called_once_with("snapshot") + + +def test_state_command_prints_source_and_encoder_details(): + client = _make_client() + client.state = MagicMock(return_value=_make_state()) + + result = CliRunner().invoke(client.cli(), ["state"]) + + assert result.exit_code == 0 + assert "Online: True" in result.output + assert "Resolution: 1280x720" in result.output + assert "FPS: 29/30" in result.output + assert "Encoder: CPU (quality: 85)" in result.output + + +def test_snapshot_command_saves_snapshot_to_requested_path(): + client = _make_client() + image = MagicMock() + client.snapshot = MagicMock(return_value=image) + + result = CliRunner().invoke(client.cli(), ["snapshot", "--output", "frame.jpg"]) + + assert result.exit_code == 0 + image.save.assert_called_once_with("frame.jpg") + assert "Saved snapshot to frame.jpg" in result.output + + +def test_stream_command_registers_routes_and_starts_server(): + client = _make_client() + client.call_async = AsyncMock(return_value=base64.b64encode(b"jpeg-data").decode("ascii")) + + captured = {} + proxied_response = object() + + with ( + patch( + "jumpstarter_driver_ustreamer.client._run_server", + side_effect=lambda client_arg, app, port, browser: captured.update( + {"client": client_arg, "app": app, "port": port, "browser": browser} + ), + ), + patch( + "jumpstarter_driver_ustreamer.client._proxy_mjpeg_stream", + new=AsyncMock(return_value=proxied_response), + ) as mock_proxy, + ): + result = CliRunner().invoke(client.cli(), ["stream", "--port", "1234", "--no-browser"]) + + assert result.exit_code == 0 + assert captured["client"] is client + assert captured["port"] == 1234 + assert captured["browser"] is False + + async def exercise_routes(): + app = captured["app"] + index_handler = _get_route_handler(app, "/") + snapshot_handler = _get_route_handler(app, "/snapshot") + stream_handler = _get_route_handler(app, "/stream") + + index_response = await index_handler(object()) + assert index_response.text == LANDING_PAGE + + snapshot_response = await snapshot_handler(object()) + assert snapshot_response.body == b"jpeg-data" + assert snapshot_response.content_type == "image/jpeg" + + request = object() + response = await stream_handler(request) + assert response is proxied_response + mock_proxy.assert_awaited_once_with(client, request) + + anyio.run(exercise_routes) + + client.call_async.assert_awaited_once_with("snapshot") + + +def test_run_server_uses_public_site_port_and_cleans_up(): + runner = SimpleNamespace(setup=AsyncMock(), cleanup=AsyncMock(), addresses=[("127.0.0.1", 59172)]) + site = SimpleNamespace(start=AsyncMock()) + client = SimpleNamespace(portal=SimpleNamespace(call=lambda fn, *args: anyio.run(fn, *args))) + + async def raise_keyboard_interrupt(): + raise KeyboardInterrupt + + with ( + patch("jumpstarter_driver_ustreamer.client.web.AppRunner", return_value=runner), + patch("jumpstarter_driver_ustreamer.client.web.TCPSite", return_value=site), + patch( + "jumpstarter_driver_ustreamer.client.move_on_after", + side_effect=lambda *args, **kwargs: contextlib.nullcontext(), + ), + patch("anyio.sleep_forever", new=raise_keyboard_interrupt), + patch("jumpstarter_driver_ustreamer.client.webbrowser.open") as mock_open, + patch("jumpstarter_driver_ustreamer.client.click.echo") as mock_echo, + ): + _run_server(client, object(), 0, True) + + runner.setup.assert_awaited_once() + site.start.assert_awaited_once() + runner.cleanup.assert_awaited_once() + mock_open.assert_called_once_with("http://127.0.0.1:59172") + mock_echo.assert_any_call("Video stream available at: http://127.0.0.1:59172") + mock_echo.assert_any_call("Snapshot endpoint: http://127.0.0.1:59172/snapshot") + mock_echo.assert_any_call("Press Ctrl+C to stop.") + mock_echo.assert_any_call("\nStopping video server.") + + +def test_run_server_propagates_startup_errors(): + runner = SimpleNamespace(setup=AsyncMock(), cleanup=AsyncMock(), addresses=[]) + site = SimpleNamespace(start=AsyncMock(side_effect=OSError("port in use"))) + client = SimpleNamespace(portal=SimpleNamespace(call=lambda fn, *args: anyio.run(fn, *args))) + + with ( + patch("jumpstarter_driver_ustreamer.client.web.AppRunner", return_value=runner), + patch("jumpstarter_driver_ustreamer.client.web.TCPSite", return_value=site), + patch( + "jumpstarter_driver_ustreamer.client.move_on_after", + side_effect=lambda *args, **kwargs: contextlib.nullcontext(), + ), + patch("jumpstarter_driver_ustreamer.client.click.echo") as mock_echo, + ): + with pytest.raises(OSError, match="port in use"): + _run_server(client, object(), 0, False) + + runner.setup.assert_awaited_once() + site.start.assert_awaited_once() + runner.cleanup.assert_awaited_once() + mock_echo.assert_not_called() + + +def test_run_server_raises_when_no_bound_address_is_reported(): + runner = SimpleNamespace(setup=AsyncMock(), cleanup=AsyncMock(), addresses=[]) + site = SimpleNamespace(start=AsyncMock()) + client = SimpleNamespace(portal=SimpleNamespace(call=lambda fn, *args: anyio.run(fn, *args))) + + with ( + patch("jumpstarter_driver_ustreamer.client.web.AppRunner", return_value=runner), + patch("jumpstarter_driver_ustreamer.client.web.TCPSite", return_value=site), + patch( + "jumpstarter_driver_ustreamer.client.move_on_after", + side_effect=lambda *args, **kwargs: contextlib.nullcontext(), + ), + ): + with pytest.raises(RuntimeError, match="without a bound address"): + _run_server(client, object(), 0, False) + + runner.setup.assert_awaited_once() + site.start.assert_awaited_once() + runner.cleanup.assert_awaited_once() + + +@pytest.mark.anyio +async def test_proxy_mjpeg_stream_forwards_headers_and_body_chunks(): + tunnel = _FakeTunnel( + [ + ( + b"HTTP/1.1 200 OK\r\n" + b"Content-Type: multipart/x-mixed-replace; boundary=frame\r\n" + b"\r\n" + b"--frame-1" + ), + b"--frame-2", + anyio.EndOfStream(), + ] + ) + client = SimpleNamespace(stream_async=lambda method: _FakeStreamContext(tunnel)) + response = _FakeStreamResponse() + request = object() + + with patch("jumpstarter_driver_ustreamer.client.web.StreamResponse", return_value=response): + result = await _proxy_mjpeg_stream(client, request) + + assert result is response + assert tunnel.sent == [b"GET /stream HTTP/1.1\r\nHost: localhost\r\n\r\n"] + assert response.content_type == "multipart/x-mixed-replace; boundary=frame" + assert response.prepared_request is request + assert response.writes == [b"--frame-1", b"--frame-2"] diff --git a/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver.py b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver.py index 4a6b3d11f..c737b0035 100644 --- a/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver.py +++ b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver.py @@ -1,5 +1,8 @@ +import ctypes +import signal import sys from base64 import b64encode +from collections.abc import Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field from pathlib import Path @@ -13,6 +16,8 @@ from .common import UStreamerState from jumpstarter.driver import Driver, export, exportstream +_IS_LINUX = sys.platform.startswith("linux") + def find_ustreamer(): executable = which("ustreamer") @@ -23,6 +28,30 @@ def find_ustreamer(): return executable +def _get_preexec_fn() -> Callable[[], None] | None: + """Get platform-specific preexec_fn for the ustreamer subprocess. + + On Linux, returns a function that sets PR_SET_PDEATHSIG to SIGTERM, + ensuring ustreamer receives SIGTERM when the parent process dies. + This works even if the parent is killed with SIGKILL. + + On other platforms, returns None. + """ + if not _IS_LINUX: + return None + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + PR_SET_PDEATHSIG = 1 + + def set_pdeathsig(): + """Set parent death signal to SIGTERM via prctl.""" + if libc.prctl(PR_SET_PDEATHSIG, signal.SIGTERM, 0, 0, 0) != 0: + errno = ctypes.get_errno() + raise OSError(errno, "prctl(PR_SET_PDEATHSIG) failed") + + return set_pdeathsig + + @dataclass(kw_only=True) class UStreamer(Driver): executable: str = field(default_factory=find_ustreamer) @@ -46,7 +75,12 @@ def __post_init__(self): cmdline += ["--unix", self.socketp] - self.process = Popen(cmdline, stdout=sys.stdout, stderr=sys.stderr) + self.process = Popen( + cmdline, + stdout=sys.stdout, + stderr=sys.stderr, + preexec_fn=_get_preexec_fn(), + ) def close(self): self.process.terminate() diff --git a/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver_test.py b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver_test.py index aae9e9c38..3292592bc 100644 --- a/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver_test.py +++ b/python/packages/jumpstarter-driver-ustreamer/jumpstarter_driver_ustreamer/driver_test.py @@ -1,6 +1,9 @@ +import signal +from unittest.mock import MagicMock, patch + import pytest -from jumpstarter_driver_ustreamer.driver import UStreamer +from jumpstarter_driver_ustreamer.driver import UStreamer, _get_preexec_fn from jumpstarter.common.utils import serve @@ -9,8 +12,63 @@ def test_drivers_video_ustreamer(): try: instance = UStreamer() except FileNotFoundError: - pytest.skip("ustreamer not available") # ty: ignore[call-non-callable] + pytest.skip("ustreamer not available") # ty: ignore[call-non-callable] with serve(instance) as client: assert client.state().ok _ = client.snapshot() + + +def test_get_preexec_fn_non_linux(): + with patch("jumpstarter_driver_ustreamer.driver._IS_LINUX", False): + assert _get_preexec_fn() is None + + +def test_get_preexec_fn_sets_pdeathsig_on_linux(): + with ( + patch("jumpstarter_driver_ustreamer.driver._IS_LINUX", True), + patch("jumpstarter_driver_ustreamer.driver.ctypes.CDLL") as mock_cdll, + ): + libc = mock_cdll.return_value + libc.prctl.return_value = 0 + + preexec_fn = _get_preexec_fn() + + assert preexec_fn is not None + assert callable(preexec_fn) + preexec_fn() + + mock_cdll.assert_called_once_with("libc.so.6", use_errno=True) + libc.prctl.assert_called_once_with(1, signal.SIGTERM, 0, 0, 0) + + +def test_get_preexec_fn_raises_when_prctl_fails(): + with ( + patch("jumpstarter_driver_ustreamer.driver._IS_LINUX", True), + patch("jumpstarter_driver_ustreamer.driver.ctypes.CDLL") as mock_cdll, + patch("jumpstarter_driver_ustreamer.driver.ctypes.get_errno", return_value=22), + ): + mock_cdll.return_value.prctl.return_value = 1 + + preexec_fn = _get_preexec_fn() + + assert preexec_fn is not None + with pytest.raises(OSError, match="prctl\\(PR_SET_PDEATHSIG\\) failed") as exc_info: + preexec_fn() + + exc = exc_info.value + assert isinstance(exc, OSError) + assert exc.errno == 22 + + +def test_ustreamer_passes_preexec_fn_to_popen(): + mock_proc = MagicMock() + with ( + patch("jumpstarter_driver_ustreamer.driver.Popen", return_value=mock_proc) as mock_popen, + patch("jumpstarter_driver_ustreamer.driver._get_preexec_fn", return_value=MagicMock(name="preexec")) as mock_fn, + ): + instance = UStreamer(executable="/usr/bin/ustreamer", args={"device": "/dev/video0"}) + mock_popen.assert_called_once() + assert mock_popen.call_args.kwargs["preexec_fn"] is mock_fn.return_value + instance.close() + mock_proc.terminate.assert_called_once()