"""The streaming proxy: HTTP range serving, and the safety rails. Ported from the standalone `proxy/test_range.py` and `proxy/test_limits.py` so that one `pytest` run covers everything. These now drive the real `make_handler(mgr, …)` rather than the PoC's `make_handler(path, done)`, which means the routing, the video-id validation and the wait-for-complete logic are exercised too. No network, no yt-dlp, no ffmpeg: the pipeline runner is stubbed and the output file is synthetic. """ from __future__ import annotations import http.client import importlib.util import json import os import threading import time from http.server import ThreadingHTTPServer from pathlib import Path import pytest PROXY_PATH = Path(__file__).resolve().parent.parent / "proxy" / "ytstream_proxy.py" def _load_proxy(): spec = importlib.util.spec_from_file_location("ytstream_proxy", PROXY_PATH) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module proxy = _load_proxy() # Position-identifiable body, so a wrong range returns visibly wrong bytes rather # than merely the wrong length. BODY = bytes(range(256)) * 400 TOTAL = len(BODY) VIDEO_ID = "dQw4w9WgXcQ" class StubSession: """A Session whose output file already exists.""" def __init__(self, path: Path, *, finished: bool = True): self.video_id = VIDEO_ID self.out_path = str(path) self.final = threading.Event() self.failed = None self.readers = 0 self.last_used = time.time() if finished: self.final.set() @property def complete(self): return self.final.is_set() and not self.failed def size(self): try: return os.path.getsize(self.out_path) except OSError: return 0 class StubManager: def __init__(self, session, *, growing=False): self.session = session self.growing = growing self.lock = threading.Lock() def get(self, video_id): return self.session, None def status(self): return {"mode": "stub", "cache_used_gb": 0.0} @pytest.fixture() def complete_server(tmp_path): """A server over a finished file, with ranges honoured.""" path = tmp_path / "out.mp4" path.write_bytes(BODY) session = StubSession(path, finished=True) yield from _serve(StubManager(session)) @pytest.fixture() def growing_server(tmp_path): """A server over a file still being written.""" path = tmp_path / "out.mp4" path.write_bytes(BODY[:1000]) session = StubSession(path, finished=False) manager = StubManager(session, growing=True) def finish(): time.sleep(0.2) with open(path, "ab") as handle: handle.write(BODY[1000:]) session.final.set() threading.Thread(target=finish, daemon=True).start() yield from _serve(manager) def _serve(manager, wait_timeout=30, grace=proxy.FIRST_BYTE_GRACE): server = ThreadingHTTPServer( ("127.0.0.1", 0), proxy.make_handler(manager, wait_timeout, grace)) threading.Thread(target=server.serve_forever, daemon=True).start() try: yield server.server_address[1] finally: server.shutdown() server.server_close() def request(port, headers=None, method="GET", path=f"/watch/{VIDEO_ID}"): conn = http.client.HTTPConnection("127.0.0.1", port, timeout=20) conn.request(method, path, headers=headers or {}) response = conn.getresponse() body = response.read() head = dict(response.getheaders()) conn.close() return response.status, head, body # ------------------------------------------------------------- complete file def test_bounded_range(complete_server): status, head, body = request(complete_server, {"Range": "bytes=1000-1999"}) assert status == 206 assert len(body) == 1000 assert body == BODY[1000:2000] assert head["Content-Range"] == f"bytes 1000-1999/{TOTAL}" assert head["Content-Length"] == "1000" def test_open_ended_range(complete_server): status, head, body = request(complete_server, {"Range": "bytes=102000-"}) assert status == 206 assert body == BODY[102000:] assert head["Content-Range"] == f"bytes 102000-{TOTAL - 1}/{TOTAL}" def test_suffix_range(complete_server): """bytes=-N means the final N bytes, not the first N.""" status, head, body = request(complete_server, {"Range": "bytes=-100"}) assert status == 206 assert len(body) == 100 assert body == BODY[-100:] assert head["Content-Range"] == f"bytes {TOTAL - 100}-{TOTAL - 1}/{TOTAL}" def test_single_byte_range(complete_server): status, _, body = request(complete_server, {"Range": "bytes=5-5"}) assert status == 206 assert body == BODY[5:6] def test_end_beyond_eof_is_clamped(complete_server): status, head, body = request(complete_server, {"Range": f"bytes=102000-{TOTAL + 5000}"}) assert status == 206 assert body == BODY[102000:] assert head["Content-Range"] == f"bytes 102000-{TOTAL - 1}/{TOTAL}" def test_range_past_eof_is_416(complete_server): """And it must carry Content-Range: bytes */total, or clients retry forever.""" status, head, _ = request(complete_server, {"Range": f"bytes={TOTAL + 10}-"}) assert status == 416 assert head["Content-Range"] == f"bytes */{TOTAL}" def test_no_range_serves_the_whole_body(complete_server): status, head, body = request(complete_server) assert status == 200 assert body == BODY assert head["Accept-Ranges"] == "bytes" def test_multi_range_falls_back_to_the_whole_body(complete_server): """Legal, and beats mis-serving one part of a multipart response.""" status, _, body = request(complete_server, {"Range": "bytes=0-99,200-299"}) assert status == 200 assert body == BODY @pytest.mark.parametrize("header", ["bytes=", "bytes=abc-def", "items=0-99", "0-99"]) def test_unparseable_range_serves_the_whole_body(complete_server, header): status, _, body = request(complete_server, {"Range": header}) assert status == 200 assert len(body) == TOTAL def test_head_returns_headers_and_no_body(complete_server): status, head, body = request(complete_server, method="HEAD") assert status == 200 assert body == b"" assert head["Content-Length"] == str(TOTAL) # ------------------------------------------------------------- growing file def test_growing_file_is_chunked_and_tracks_to_eof(growing_server): """Ranges cannot be honoured mid-write: there is no reliable time-to-byte mapping into a fragmented MP4 yet, so it becomes a non-seekable stream.""" status, head, body = request(growing_server, {"Range": "bytes=500-999"}) assert status == 200 assert head.get("Transfer-Encoding") == "chunked" assert body == BODY # ------------------------------------------------- time-to-first-byte grace # # Waiting for a complete mux is the only way to get a correct duration and # working seeks, but the wait grows with the video's length and sends nothing at # all -- not even headers -- while it lasts. A 46-minute upload blocked for 79s # and Jellyfin reported a playback error on every first play of a long video. # These tests pin the compromise: wait, but only for a bounded grace. def _partial(tmp_path, *, finishes_after): """A half-written file that completes after `finishes_after` seconds.""" path = tmp_path / "out.mp4" path.write_bytes(BODY[:1000]) session = StubSession(path, finished=False) def finish(): time.sleep(finishes_after) with open(path, "ab") as handle: handle.write(BODY[1000:]) session.final.set() threading.Thread(target=finish, daemon=True).start() return session def test_a_slow_mux_streams_instead_of_blocking(tmp_path): """The bug: this used to block for --wait-timeout with the socket silent.""" session = _partial(tmp_path, finishes_after=1.0) server = _serve(StubManager(session), wait_timeout=30, grace=0.2) port = next(server) try: started = time.monotonic() status, head, body = request(port, {"Range": "bytes=0-"}) elapsed = time.monotonic() - started assert status == 200 assert head.get("Transfer-Encoding") == "chunked" assert body == BODY # The whole point: it returned on the mux's schedule, not the timeout's. assert elapsed < 10, f"blocked {elapsed:.1f}s -- grace not applied" finally: server.close() def test_a_mux_that_lands_inside_the_grace_keeps_ranges_and_seeking(tmp_path): """The fast path must survive: a short video still gets a real 206.""" session = _partial(tmp_path, finishes_after=0.2) server = _serve(StubManager(session), wait_timeout=30, grace=10) port = next(server) try: status, head, body = request(port, {"Range": "bytes=100-199"}) assert status == 206 assert head["Content-Range"] == f"bytes 100-199/{TOTAL}" assert head.get("Transfer-Encoding") is None assert body == BODY[100:200] finally: server.close() def test_grace_at_the_wait_timeout_restores_strict_finished_file_only(tmp_path): """The old behaviour stays reachable, for a caller that really wants it.""" session = _partial(tmp_path, finishes_after=60) server = _serve(StubManager(session), wait_timeout=0.3, grace=0.3) port = next(server) try: status, _, _ = request(port) assert status == 504 finally: server.close() def test_a_failed_producer_with_bytes_is_served_rather_than_erroring(tmp_path): """Most of a video beats none of it -- except in strict mode, where the caller asked for a good file and must be told it cannot have one.""" path = tmp_path / "out.mp4" path.write_bytes(BODY) session = StubSession(path, finished=True) session.failed = "ffmpeg exited 1" lenient = _serve(StubManager(session), wait_timeout=30, grace=1) port = next(lenient) try: assert request(port)[0] == 200 finally: lenient.close() strict = _serve(StubManager(session), wait_timeout=1, grace=1) port = next(strict) try: assert request(port)[0] == 502 finally: strict.close() def test_a_failed_producer_with_no_bytes_is_always_502(tmp_path): path = tmp_path / "out.mp4" path.write_bytes(b"") session = StubSession(path, finished=True) session.failed = "yt-dlp[video] exited 1" server = _serve(StubManager(session), wait_timeout=30, grace=1) port = next(server) try: assert request(port)[0] == 502 finally: server.close() def test_growing_flag_means_no_wait_at_all(tmp_path): """--growing is a grace of zero, and must not be overridden by the default.""" session = _partial(tmp_path, finishes_after=0.5) manager = StubManager(session, growing=True) server = _serve(manager, wait_timeout=30, grace=proxy.FIRST_BYTE_GRACE) port = next(server) try: started = time.monotonic() status, head, _ = request(port) assert status == 200 assert head.get("Transfer-Encoding") == "chunked" # Would have waited the full 12s default grace if growing were ignored. assert time.monotonic() - started < 5 finally: server.close() def test_status_reports_the_serving_mode_and_grace(manager_factory): """/healthz has to show this: it is the difference between 'plays' and 'playback error', and it is otherwise invisible without the unit file.""" manager = manager_factory() status = manager.status() assert status["mode"] == "wait-then-stream" assert status["first_byte_grace_s"] == proxy.FIRST_BYTE_GRACE growing = manager_factory(growing=True) assert growing.status()["mode"] == "growing" assert growing.status()["first_byte_grace_s"] == 0.0 # ------------------------------------------------------------------- routing def test_unknown_path_is_404(complete_server): status, _, _ = request(complete_server, path="/nope") assert status == 404 def test_malformed_video_id_is_400(complete_server): status, _, _ = request(complete_server, path="/watch/short") assert status == 400 @pytest.mark.parametrize("bad", ["../../etc/passwd", "abcdefghij", "abcdefghijkl"]) def test_video_id_must_be_exactly_eleven_safe_chars(complete_server, bad): status, _, _ = request(complete_server, path=f"/watch/{bad}") assert status in (400, 404) def test_healthz_returns_json(complete_server): status, head, body = request(complete_server, path="/healthz") assert status == 200 assert head["Content-Type"] == "application/json" assert json.loads(body)["mode"] == "stub" # ------------------------------------------------- safety rails on the manager def vid(n): return f"vid{n:08d}" @pytest.fixture() def manager_factory(tmp_path): made = [] def build(**kwargs): options = dict(max_pipelines=99, cache_bytes=10 ** 12, no_fetch=False, growing=False, max_retries=0, max_starts=3, starts_window=3600.0) options.update(kwargs) work = tmp_path / f"work{len(made)}" work.mkdir() class Recording(proxy.Manager): """Records starts instead of spawning a pipeline.""" def __init__(self, *args, **kw): super().__init__(*args, **kw) self.ran = [] def _run(self, session): self.ran.append(session.video_id) session.final.set() # completes instantly, never stays active manager = Recording(str(work), **options) made.append(manager) return manager return build def test_cold_start_budget_is_enforced(manager_factory): manager = manager_factory(max_starts=3) results = [manager.get(vid(i)) for i in range(5)] started = [r for r in results if r[0] is not None] refused = [r for r in results if r[0] is None] assert len(started) == 3 assert len(refused) == 2 assert all("cold-start budget" in r[1] for r in refused) assert manager.counters["refused_ratelimit"] == 2 assert len(manager.ran) == 3 def test_cache_hits_are_never_rate_limited(manager_factory): """Re-watching an already-fetched video must keep working with the budget spent, or a metadata refresh would break normal playback for an hour.""" manager = manager_factory(max_starts=3) for i in range(5): manager.get(vid(i)) session, refusal = manager.get(vid(0)) assert session is not None and refusal is None assert manager.counters["reused"] >= 1 assert len(manager.ran) == 3 def test_start_budget_window_prunes(manager_factory): manager = manager_factory(max_starts=2) manager.get(vid(10)) manager.get(vid(11)) assert manager.get(vid(12))[0] is None # Age the recorded starts past the window. manager.start_log = type(manager.start_log)(t - 3601 for t in manager.start_log) assert manager.get(vid(13))[0] is not None assert len(manager.start_log) == 1 def test_concurrency_cap(tmp_path): class Blocking(proxy.Manager): """Pipelines that never finish, so sessions stay active.""" def _run(self, session): pass work = tmp_path / "blocking" work.mkdir() manager = Blocking(str(work), max_pipelines=2, cache_bytes=10 ** 12, no_fetch=False, growing=False, max_retries=0, max_starts=99, starts_window=3600.0) results = [manager.get(vid(20 + i)) for i in range(4)] ok = [r for r in results if r[0] is not None] refused = [r for r in results if r[0] is None] assert len(ok) == 2 assert len(refused) == 2 assert all("pipeline cap" in r[1] for r in refused) # A busy refusal must not consume start budget, or a burst of concurrent # requests would exhaust the hourly allowance without fetching anything. assert manager.counters["refused_busy"] == 2 assert len(manager.start_log) == 2 def test_startup_clears_untracked_cache(tmp_path): """The work root is a tmpfs and the session map is memory-only, so anything left by a previous run is unreachable *and* unevictable -- it would leak RAM until the next reboot.""" work = tmp_path / "work" (work / "abcdefghijk").mkdir(parents=True) (work / "abcdefghijk" / "out.mp4").write_bytes(b"x" * 5000) (work / "bcdefghijkl").mkdir() (work / "bcdefghijkl" / "out.mp4").write_bytes(b"y" * 3000) (work / "loose.txt").write_text("not a session") reclaimed = proxy.reset_work_root(str(work)) assert reclaimed == 8000 assert not (work / "abcdefghijk").exists() assert not (work / "bcdefghijkl").exists() # A stray file is not a session directory and is left alone. assert (work / "loose.txt").exists() def test_startup_on_a_clean_work_root_is_a_no_op(tmp_path): work = tmp_path / "empty" work.mkdir() assert proxy.reset_work_root(str(work)) == 0 def test_no_fetch_mode_refuses_everything(manager_factory): manager = manager_factory(no_fetch=True, max_starts=99) session, refusal = manager.get(vid(30)) assert session is None assert "no-fetch" in refusal assert manager.ran == [] assert len(manager.start_log) == 0