mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 20:15:10 +00:00
imSp4rky
3027592e31
Vaults.add_keys and kv add skip null content keys instead of failing the batch; is_null_key moves to core.vault. A manifest snapshot restore has one 32 MiB inflate budget. Tests reset the speed limiter and config between runs.
141 lines
4.2 KiB
Python
141 lines
4.2 KiB
Python
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()."""
|
|
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
|
|
|
|
|
|
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()
|