ytstream/tests/test_proxy.py

"""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