unshackle/tests/downloaders/test_multiprocess.py

"""Multiprocess segment-download invariants.

Pins the spawned-child fan-out path in ``requests()``/``download_multiprocess``:

- Strided assignment plus the ``index_offset``/``index_stride`` filename mapping puts every
  segment's bytes in the correctly named file (a broken stride mapping misfiles content).
- The parent normalizes child progress to segment granularity: it drops child ``advance``
  and ``total``/speed events and emits exactly one ``advance=1`` per ``file_downloaded``, so
  the advances summed by a caller equal the segment count.
- A child that dies (its worker sends ``__mp_error__``) surfaces as a ``RuntimeError`` from
  the generator without hanging.
- A process-global ``DOWNLOAD_CANCELLED`` set mid-batch terminates the spawned children (which
  carry their own fresh flag and can't see it) and the teardown join reaps them promptly.

Children are spawned and re-import the module, so monkeypatched module attributes do NOT reach
them; these tests drive real child processes against a localhost server on 127.0.0.1 (reachable
from the children) and avoid any child-side patching. Segments are a few hundred bytes so the
real spawns stay fast.
"""

import importlib
import multiprocessing
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import pytest

dl = importlib.import_module("unshackle.core.downloaders.requests")

SEG_REPEAT = 32  # 8-digit index * 32 = 256-byte body, distinct per index


def seg_body(index: int) -> bytes:
    """Segment body that encodes its own index, so misfiled striding is detectable."""
    return (f"{index:08d}" * SEG_REPEAT).encode()


class _Handler(BaseHTTPRequestHandler):
    def do_GET(self) -> None:
        server = self.server
        delay = getattr(server, "delay", 0.0)
        if delay:
            time.sleep(delay)
        try:
            index = int(self.path.rsplit("/", 1)[1].split(".")[0])
        except ValueError:
            self.send_response(404)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        body = seg_body(index)
        self.send_response(200)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *args):  # keep the test output quiet
        pass


class _Server(ThreadingHTTPServer):
    daemon_threads = True

    def handle_error(self, request, client_address):  # a cancelled child closes its socket mid-write
        pass


@pytest.fixture
def server():
    srv = _Server(("127.0.0.1", 0), _Handler)
    srv.delay = 0.0
    thread = threading.Thread(target=srv.serve_forever, daemon=True)
    thread.start()
    yield srv
    srv.shutdown()
    thread.join(timeout=5)


@pytest.fixture(autouse=True)
def clear_cancel():
    # spawned children carry their own flag, but the parent loop reads this process's flag;
    # keep it clean so cancel state never leaks between tests
    dl.DOWNLOAD_CANCELLED.clear()
    yield
    dl.DOWNLOAD_CANCELLED.clear()


def url(srv, path):
    return f"http://127.0.0.1:{srv.server_address[1]}{path}"


def test_multiprocess_stride_downloads_all_segments_correctly(server, tmp_path):
    # >= MP_MIN_SEGMENTS with processes=2 engages the fan-out via the public entry point
    n = dl.MP_MIN_SEGMENTS
    urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)]

    total_advance = 0
    for event in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
        total_advance += event.get("advance", 0)

    files = sorted(tmp_path.glob("seg_*.bin"))
    assert len(files) == n
    for f in files:
        index = int(f.stem.split("_")[1])
        assert f.read_bytes() == seg_body(index), f"segment {index} holds another index's bytes (stride mapping bug)"
    # parent drops child advances and emits one advance=1 per file_downloaded
    assert total_advance == n


def test_multiprocess_child_crash_raises(server, tmp_path):
    # download_multiprocess has no MP_MIN_SEGMENTS gate, so drive it directly with 2 urls.
    # a dict url without "url" raises KeyError while requests() builds the save path (before
    # any retry loop), so the child sends __mp_error__ fast; the sibling downloads normally.
    good = {"url": url(server, "/seg/0.bin")}
    bad = {"nourl": True}
    with pytest.raises(RuntimeError):
        for _ in dl.download_multiprocess(
            urls=[good, bad],
            output_dir=tmp_path,
            filename="seg_{i:04}.bin",
            headers=None,
            cookies=None,
            proxy=None,
            max_workers=2,
            adaptive=False,
            spec={"kind": "none"},
            processes=2,
            debug_logger=None,
        ):
            pass


def test_multiprocess_cancel_terminates_children(server, tmp_path):
    server.delay = 0.5  # slow segments so cancel lands well before the batch could finish
    n = dl.MP_MIN_SEGMENTS
    urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)]

    baseline = len(multiprocessing.active_children())
    start = time.time()
    progressed = 0
    for event in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
        if event.get("advance") or event.get("written"):
            progressed += 1
            if progressed >= 2:
                dl.DOWNLOAD_CANCELLED.set()  # sibling-track cancel: parent must tear the children down
    elapsed = time.time() - start

    # cancel stopped the work: nowhere near all segments finished, and no hang
    assert len(list(tmp_path.glob("seg_*.bin"))) < n
    assert elapsed < 15
    # the finally-block join reaps the spawned children promptly
    deadline = time.time() + 5
    while time.time() < deadline and len(multiprocessing.active_children()) > baseline:
        time.sleep(0.05)
    assert len(multiprocessing.active_children()) <= baseline


def test_speed_limit_forces_single_process(server, tmp_path, monkeypatch):
    """A set speed limit must keep the batch in-process: spawned children re-import the
    module and would run with no TokenBucket, silently ignoring the configured cap."""

    def no_mp(**kwargs):
        raise AssertionError("multiprocess fan-out engaged despite a speed limit")

    monkeypatch.setattr(dl, "download_multiprocess", no_mp)
    n = dl.MP_MIN_SEGMENTS
    urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)]
    dl.set_speed_limit(1_000_000_000)  # far above the tiny batch, so the test is not slowed
    try:
        for _ in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
            pass
    finally:
        dl.set_speed_limit(None)

    files = sorted(tmp_path.glob("seg_*.bin"))
    assert len(files) == n
    for f in files:
        assert f.read_bytes() == seg_body(int(f.stem.split("_")[1]))