unshackle/tests/downloaders/test_speed_limit.py

import importlib
import inspect
import threading
import time
from types import SimpleNamespace

import pytest

from unshackle.core.constants import DOWNLOAD_CANCELLED
from unshackle.core.downloaders.requests import TokenBucket, format_speed, parse_speed_limit, set_speed_limit

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


@pytest.fixture(autouse=True)
def clean_cancel():
    # a failing batch sets the process-global cancel, which TokenBucket.consume waits on;
    # keep tests independent
    DOWNLOAD_CANCELLED.clear()
    yield
    DOWNLOAD_CANCELLED.clear()


@pytest.mark.parametrize(
    ("value", "expected"),
    [
        ("500k", 500_000),
        ("5M", 5_000_000),
        ("10MB", 10_000_000),
        ("10mb/s", 10_000_000),
        ("1.5G", 1_500_000_000),
        ("1048576", 1_048_576),
        (5_000_000, 5_000_000.0),
        (None, None),
        ("off", None),
        ("none", None),
        ("unlimited", None),
        ("0", None),
        (0, None),
    ],
)
def test_parse_speed_limit(value, expected):
    assert parse_speed_limit(value) == expected


@pytest.mark.parametrize(
    "value",
    ["fast", "10TB", "10MiB", "-5M", "M", "1..5M", ".", "1.2.3", float("inf"), float("nan"), -5, -1.5],
)
def test_parse_speed_limit_rejects(value):
    with pytest.raises(ValueError):
        parse_speed_limit(value)


@pytest.mark.parametrize(
    ("rate", "expected"),
    [
        (500, "500 bytes/s"),
        (512_000, "512.0 kB/s"),
        (39 * 1024**2, "40.9 MB/s"),
        (1.5 * 1024**3, "1.6 GB/s"),
    ],
)
def test_format_speed(rate, expected):
    assert format_speed(rate) == expected


def test_result_accepts_a_missing_speed_limit():
    """The API/serve path calls dl.result() by keyword without speed_limit."""
    from unshackle.commands.dl import dl

    assert inspect.signature(dl.result).parameters["speed_limit"].default is None


def test_bucket_holds_aggregate_rate_across_threads():
    rate = 10_000_000
    bucket = TokenBucket(rate)
    chunk = 100_000
    totals = [0, 0, 0, 0]
    start = time.monotonic()
    stop = start + 1.0

    def worker(i: int) -> None:
        while time.monotonic() < stop:
            bucket.consume(chunk)
            totals[i] += chunk

    threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    achieved = sum(totals) / (time.monotonic() - start)
    # exceeding the limit is a real failure; undershoot is scheduler noise on a loaded box, so no lower bound
    assert achieved <= rate * 1.1


def test_bucket_refills_at_the_rate(monkeypatch):
    """On a simulated clock the throughput is exact, so a slow refill fails like a fast one."""
    clock = [0.0]
    monkeypatch.setattr(requests_downloader, "time", SimpleNamespace(monotonic=lambda: clock[0]))
    monkeypatch.setattr(
        requests_downloader, "DOWNLOAD_CANCELLED", SimpleNamespace(wait=lambda s: clock.__setitem__(0, clock[0] + s))
    )
    bucket = TokenBucket(1_000)
    for _ in range(50):
        bucket.consume(100)
    assert clock[0] == pytest.approx(5.0)


def test_bucket_chunk_larger_than_capacity_does_not_deadlock(monkeypatch):
    # capture the debt sleep instead of timing it: Event.wait can wake early spuriously,
    # so a wall-clock lower bound flakes
    waits: list[float] = []
    monkeypatch.setattr(requests_downloader, "DOWNLOAD_CANCELLED", SimpleNamespace(wait=waits.append))
    bucket = TokenBucket(1_000_000)
    bucket.consume(1_100_000)
    assert len(waits) == 1
    assert 1.0 < waits[0] < 2.0


def test_locked_limit_ignores_unlocked_calls():
    """serve's global_speed_limit must survive the per-job set in dl.result()."""
    try:
        set_speed_limit(5_000_000, lock=True)
        set_speed_limit(None)
        set_speed_limit(1_000_000)
        limiter = requests_downloader.speed_limiter
        assert limiter is not None and limiter.rate == 5_000_000
    finally:
        requests_downloader.speed_limit_locked = False
        requests_downloader.speed_limiter = None


def test_bucket_wait_aborts_on_cancellation():
    """Debt sleeps grow with thread count; cancelling must not wait them out."""
    bucket = TokenBucket(1_000)
    DOWNLOAD_CANCELLED.set()
    try:
        start = time.monotonic()
        bucket.consume(1_000_000)
        assert time.monotonic() - start < 1.0
    finally:
        DOWNLOAD_CANCELLED.clear()