unshackle/tests/core/test_adaptive_workers.py

"""Unit tests for the opt-in adaptive worker controller (AIMD hill-climb over segment
concurrency) and the plumbing that carries the flag into the downloader. Pure controller
tests only. No threads, sockets, or real downloads; the clock is supplied explicitly."""

from __future__ import annotations

import inspect
from functools import partial
from pathlib import Path

import pytest

from unshackle.core.downloaders.requests import (
    TAIL_BOOST_MIN_SEGMENT_SIZE,
    TAIL_BOOST_PART_SIZE,
    AdaptiveWorkerController,
    has_range_header,
    plan_tail_parts,
    requests,
    split_ranges,
    submission_window,
    tail_boost_engages,
)
from unshackle.core.tracks.track import DownloadContext

pytestmark = pytest.mark.unit

TICK = 4.0
WINDOW = 10.0
DT = 0.5


def fill_window(c: AdaptiveWorkerController, rate: float, now: float) -> float:
    """Record ``WINDOW`` seconds of samples at ``rate`` bytes/sec ending on the 0.5s grid.

    Records a sample at ``now`` and every 0.5s up to ``now + WINDOW`` so the oldest kept
    sample sits exactly ``WINDOW`` old. That is the condition ``update`` needs to start acting.
    """
    c.record_bytes(int(rate * DT), now)
    end = now + WINDOW
    while now < end - 1e-9:
        now += DT
        c.record_bytes(int(rate * DT), now)
    return now


def run_tick(c: AdaptiveWorkerController, rate: float, now: float, errors: int = 0) -> tuple[int, float]:
    """Advance one tick recording samples at ``rate``, optionally inject ``errors``, then update."""
    end = now + TICK
    while now < end - 1e-9:
        now += DT
        c.record_bytes(int(rate * DT), now)
    for _ in range(errors):
        c.record_error(now)
    return c.update(now), now


def test_starts_at_cap() -> None:
    assert AdaptiveWorkerController(cap=16).target == 16
    assert AdaptiveWorkerController(cap=4).target == 4
    assert AdaptiveWorkerController(cap=1).target == 2  # clamped up to ADAPTIVE_MIN
    assert AdaptiveWorkerController(cap=16, start=6).target == 6  # explicit start wins


def test_no_op_before_warmup() -> None:
    c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW)
    assert c.update(0.0) == 6  # arms last_tick before any samples are recorded
    c.record_bytes(500_000, 3.0)
    # a full tick has elapsed since arming, but under half a tick of samples exists -> hold
    assert c.update(4.5) == 6
    c.record_bytes(500_000, 6.5)
    # warmed up (first sample 3.0, now 8.5) -> the probe fires
    assert c.update(8.5) == 8


def test_climbs_while_speed_improves() -> None:
    c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)  # arm
    rate = 1_000_000.0
    now = fill_window(c, rate, now)
    targets = [c.update(now)]  # first evaluation: 6 -> 8
    for _ in range(6):
        rate *= 2  # sustained >10% windowed gain each tick
        target, now = run_tick(c, rate, now)
        targets.append(target)
    assert targets == sorted(targets)  # monotonic non-decreasing
    assert targets[-1] == 16  # climbed to the cap
    assert max(targets) <= 16


def test_reverts_on_plateau() -> None:
    c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    rate = 1_000_000.0
    now = fill_window(c, rate, now)
    assert c.update(now) == 8  # first increase 6 -> 8
    # hold the rate flat: <10% gain -> the tick reverts the increase to its restore point
    target, now = run_tick(c, rate, now)
    assert target == 6
    # a paying probe climbs again from the restore point
    target, now = run_tick(c, rate * 4, now)
    assert target == 8


def test_halves_and_cooldown_on_error_burst() -> None:
    c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    now = fill_window(c, 1_000_000.0, now)
    for _ in range(3):
        c.record_error(now)
    assert c.update(now) == 3  # 3 errors -> multiplicative decrease 6 // 2

    # cooldown tick holds even with improving throughput
    target, now = run_tick(c, 2_000_000.0, now)
    assert target == 3

    # cooldown consumed -> climbing resumes
    target, now = run_tick(c, 4_000_000.0, now)
    assert target == 5


def test_clamps_to_bounds_and_never_exceeds_cap() -> None:
    c = AdaptiveWorkerController(cap=4, tick=TICK, window=WINDOW)
    assert c.target == 4  # already at cap (min(6, 4))
    now = 0.0
    c.update(now)
    rate = 1_000_000.0
    now = fill_window(c, rate, now)
    seen = [c.update(now)]
    for _ in range(10):
        rate *= 2
        target, now = run_tick(c, rate, now)
        seen.append(target)
    assert all(2 <= s <= 4 for s in seen)
    assert max(seen) == 4  # never exceeds the cap


def test_error_burst_floors_at_min() -> None:
    c = AdaptiveWorkerController(cap=16, start=2, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    now = fill_window(c, 1_000_000.0, now)
    for _ in range(5):
        c.record_error(now)
    assert c.update(now) == 2  # max(ADAPTIVE_MIN, 2 // 2)


def test_tail_guard_suppresses_increase() -> None:
    # with fewer units of work than the target, the tick must hold instead of probing up
    c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)  # arm at target 6
    now = fill_window(c, 1_000_000.0, now)
    assert c.update(now, 3) == 6  # 3 < target 6 -> holds at 6 rather than climbing to 8


def test_tail_guard_none_path_unchanged() -> None:
    # without a tail count, the update still probes upward on the first evaluation
    c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    now = fill_window(c, 1_000_000.0, now)
    assert c.update(now) == 8  # first evaluation 6 -> 8


def test_tail_guard_still_halves_on_error_burst() -> None:
    # the guard suppresses growth but never blocks an error-burst decrease
    c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    now = fill_window(c, 1_000_000.0, now)
    for _ in range(3):
        c.record_error(now)
    assert c.update(now, 2) == 3  # guard active (2 < 6) but 3 errors -> 6 // 2


PER_WORKER = 1_000_000.0  # bytes/sec each worker adds, so throughput scales with the target


def run_scaling(c: AdaptiveWorkerController, now: float, seconds: float, work: int = 10_000) -> float:
    """Advance ``seconds`` of a download whose throughput scales with the live target.

    Records samples on the 0.5s grid at ``target * PER_WORKER`` bytes/sec and calls ``update``
    on every step, so the speed window mixes old and new rates the way a real run does. ``work``
    stands for the segments in flight or not yet started, kept high to hold the tail guard off.
    """
    end = now + seconds
    while now < end - 1e-9:
        now += DT
        c.record_bytes(int(c.target * PER_WORKER * DT), now)
        c.update(now, work)
    return now


def test_probe_kept_when_the_true_gain_clears_the_bar() -> None:
    c = AdaptiveWorkerController(cap=16, start=10, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    rate = 1_000_000.0
    now = fill_window(c, rate, now)
    assert c.update(now) == 12  # first evaluation probes 10 -> 12
    target, now = run_tick(c, rate * 1.2, now)
    assert target == 14  # probe held, and the controller climbs again


def test_recovers_to_the_cap_after_an_error_burst() -> None:
    c = AdaptiveWorkerController(cap=16, start=16, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    now = run_scaling(c, now, WINDOW)
    for _ in range(3):
        c.record_error(now)
    now = run_scaling(c, now, TICK)
    assert c.target == 8, "error burst must halve the target 16 -> 8"

    now = run_scaling(c, now, 60.0)
    assert c.target == 16, "controller never returned to its cap after backing off"
    now = run_scaling(c, now, 40.0)
    assert c.target == 16, "controller left the cap on a link where every worker still pays"


def test_reverts_when_the_probe_stops_paying_on_a_saturated_link() -> None:
    c = AdaptiveWorkerController(cap=16, start=8, tick=TICK, window=WINDOW)
    now = 0.0
    c.update(now)
    rate = 8 * PER_WORKER
    now = fill_window(c, rate, now)
    seen = [c.update(now, 10_000)]
    for _ in range(8):
        rate = min(c.target, 10) * PER_WORKER
        end = now + TICK
        while now < end - 1e-9:
            now += DT
            c.record_bytes(int(rate * DT), now)
        seen.append(c.update(now, 10_000))
    assert max(seen) <= 12  # one probe past the knee is the most it may hold
    assert seen[-1] <= 12


def test_submission_window_stages_a_spare_segment_for_every_worker() -> None:
    assert submission_window(16) == 32
    for target in (2, 4, 6, 16):
        assert submission_window(target) == 2 * target


def test_plan_tail_parts_full_coverage_no_gaps() -> None:
    size = 20 * 1024 * 1024
    parts = plan_tail_parts(size, spare=8)
    assert parts, "expected a split for a large segment with spare capacity"
    assert parts[0][0] == 0
    assert parts[-1][1] == size - 1
    for (_, prev_end), (start, _) in zip(parts, parts[1:]):
        assert start == prev_end + 1  # each part begins exactly where the previous one ended
    covered = sum(end - start + 1 for start, end in parts)
    assert covered == size


def test_plan_tail_parts_respects_spare_cap() -> None:
    size = 100 * 1024 * 1024  # would want ~25 parts at 4MB, but spare caps it
    parts = plan_tail_parts(size, spare=3)
    assert len(parts) <= 3
    assert parts[0][0] == 0 and parts[-1][1] == size - 1


def test_plan_tail_parts_part_size_bounded() -> None:
    size = 40 * 1024 * 1024
    parts = plan_tail_parts(size, spare=16)
    # never more parts than ~ceil(size / TAIL_BOOST_PART_SIZE)
    assert len(parts) <= -(-size // TAIL_BOOST_PART_SIZE)


def test_plan_tail_parts_skips_small_or_no_spare() -> None:
    assert plan_tail_parts(TAIL_BOOST_MIN_SEGMENT_SIZE - 1, spare=8) == []  # too small
    assert plan_tail_parts(20 * 1024 * 1024, spare=1) == []  # not enough spare workers
    assert plan_tail_parts(20 * 1024 * 1024, spare=0) == []


def test_split_ranges_covers_exactly_no_gaps_no_overlaps() -> None:
    for size in (1, 100, 4 * 1024 * 1024, 16 * 1024 * 1024 + 7, 64 * 1024 * 1024):
        parts = split_ranges(size, max_parts=8, part_target=4 * 1024 * 1024)
        assert parts[0][0] == 0 and parts[-1][1] == size - 1
        for (_, prev_end), (start, _) in zip(parts, parts[1:]):
            assert start == prev_end + 1
        assert all(start <= end for start, end in parts)
        assert sum(end - start + 1 for start, end in parts) == size


def test_split_ranges_respects_max_parts_and_target() -> None:
    assert len(split_ranges(100 * 1024 * 1024, max_parts=4, part_target=16 * 1024 * 1024)) == 4  # capped
    assert len(split_ranges(10 * 1024 * 1024, max_parts=8, part_target=16 * 1024 * 1024)) == 1  # below one target
    assert split_ranges(1, max_parts=8, part_target=16 * 1024 * 1024) == [(0, 0)]


def test_tail_boost_engages_on_idle_capacity() -> None:
    # engages at the tail: idle workers outnumber the remaining whole segments
    assert tail_boost_engages(remaining=4, pending=2, target=8)  # spare 6 >= 4
    assert tail_boost_engages(remaining=6, pending=6, target=16)  # spare 10 >= 6


def test_tail_boost_engages_when_remaining_is_a_full_stride() -> None:
    # remaining sits at 6 (one worker stride) while the target has climbed to 16, so 14 workers
    # are idle. idle capacity far exceeds the remaining segments, so the boost must engage.
    assert tail_boost_engages(remaining=6, pending=2, target=16)


def test_tail_boost_does_not_engage_early_or_saturated() -> None:
    assert not tail_boost_engages(remaining=54, pending=2, target=6)  # early: 54 > spare 4
    assert not tail_boost_engages(remaining=1, pending=8, target=8)  # saturated: spare 0
    assert not tail_boost_engages(remaining=0, pending=2, target=8)  # nothing left
    assert not tail_boost_engages(remaining=5, pending=10, target=8)  # over-subscribed: spare < 0


def test_has_range_header_detects_byte_range_slices() -> None:
    # items carrying a Range header are byte-range slices; ranged-parallel must skip them
    assert has_range_header({"url": "u", "headers": {"Range": "bytes=0-99"}})
    assert has_range_header({"url": "u", "headers": {"range": "bytes=0-99"}})  # case-insensitive
    assert not has_range_header({"url": "u", "headers": {"Authorization": "Bearer x"}})
    assert not has_range_header({"url": "u", "headers": {}})
    assert not has_range_header({"url": "u"})


def test_download_context_adaptive_default_false() -> None:
    ctx = DownloadContext(save_path=Path("x.mp4"), save_dir=Path("."), progress=partial(lambda **_: None))
    assert ctx.adaptive_workers is False


def test_requests_accepts_adaptive_kwarg() -> None:
    sig = inspect.signature(requests)
    assert "adaptive" in sig.parameters
    assert sig.parameters["adaptive"].default is False