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):
    server = ThreadingHTTPServer(("127.0.0.1", 0), proxy.make_handler(manager, 30))
    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


# ------------------------------------------------------------------- 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_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