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
from requests import Session

from unshackle.core.constants import DownloadCancelled
from unshackle.core.utils.sslciphers import SSLCiphers

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)
        path, _, query = self.path.partition("?")
        required = getattr(server, "require_query", "")
        if required and query != required:
            self.send_response(403)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        try:
            index = int(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
    srv.require_query = ""
    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
    with pytest.raises(DownloadCancelled):
        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
    for _ in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
        pass

    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]))


def test_children_carry_session_params(server, tmp_path):
    """Services keep their access token in session.params. A child that loses it gets a 403 per segment."""
    server.require_query = "at=auth-token"
    session = Session()
    session.params = {"at": "auth-token"}
    n = dl.MP_MIN_SEGMENTS
    urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)]

    for _ in dl.requests(
        urls, output_dir=tmp_path, filename="seg_{i:04}.bin", session=session, max_workers=4, processes=2
    ):
        pass

    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]))


def test_service_mounted_adapter_forces_single_process(server, tmp_path, monkeypatch):
    """A child mounts a plain adapter, so a session carrying SSLCiphers has to stay in this process."""

    def no_mp(**kwargs):
        raise AssertionError("multiprocess fan-out engaged for a session with a service-mounted adapter")

    monkeypatch.setattr(dl, "download_multiprocess", no_mp)
    session = Session()
    session.mount("https://", SSLCiphers(security_level=2))
    n = dl.MP_MIN_SEGMENTS
    urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)]

    for _ in dl.requests(
        urls, output_dir=tmp_path, filename="seg_{i:04}.bin", session=session, max_workers=4, processes=2
    ):
        pass

    assert len(list(tmp_path.glob("seg_*.bin"))) == n