mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
imSp4rky
e3c1ce3511
Completes the core underscore sweep for the test files this branch added; make_ctx/make_track/seg_urls avoid shadowing locals.
232 lines
9.7 KiB
Python
232 lines
9.7 KiB
Python
"""Regression tests for downloader resume/completeness robustness.
|
|
|
|
Each test drives the real ``download()`` generator against a local HTTP server
|
|
that injects a specific fault the fault-injection benchmark surfaced:
|
|
|
|
- truncated segment (clean FIN mid-body) must not finalize as complete
|
|
- an interrupted download must resume from the bytes actually written, not a
|
|
pre-allocated full size
|
|
- an oversized/stale resume that provokes a 416 must restart clean
|
|
- byte-range-slice segments (DASH SegmentBase / HLS EXT-X-BYTERANGE style) must
|
|
never be range-probed or tail-boosted: that would clobber their Range header
|
|
and fetch the whole resource into the segment file
|
|
|
|
All localhost + small payloads, so they run in the default (not live) suite.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import threading
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
from requests import Session
|
|
|
|
# the downloaders package re-exports a `requests` function that shadows this submodule
|
|
# attribute, so fetch the real module object from sys.modules via importlib
|
|
dl = importlib.import_module("unshackle.core.downloaders.requests")
|
|
|
|
|
|
def consume(gen) -> list:
|
|
return list(gen)
|
|
|
|
|
|
class _FaultServer:
|
|
"""Serves a fixed body with per-request fault injection and Range support."""
|
|
|
|
def __init__(self, body: bytes) -> None:
|
|
self.body = body
|
|
self.attempts = 0 # GETs seen for the single resource
|
|
self.ranges: list[Optional[str]] = [] # Range header per GET, in order
|
|
self.truncate_first = 0 # send only half the body + close, for the first N attempts
|
|
self._lock = threading.Lock()
|
|
server = self
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def log_message(self, *_a: object) -> None: # silence
|
|
pass
|
|
|
|
def do_GET(self) -> None: # noqa: N802 (stdlib name)
|
|
with server._lock:
|
|
server.attempts += 1
|
|
attempt = server.attempts
|
|
rng = self.headers.get("Range")
|
|
server.ranges.append(rng)
|
|
|
|
total = len(server.body)
|
|
start = 0
|
|
if rng and rng.startswith("bytes="):
|
|
start = int(rng.removeprefix("bytes=").split("-", 1)[0])
|
|
|
|
if start >= total: # unsatisfiable: real CDNs answer 416
|
|
self.send_response(416)
|
|
self.send_header("Content-Range", f"bytes */{total}")
|
|
self.end_headers()
|
|
return
|
|
|
|
chunk = server.body[start:]
|
|
partial = attempt <= server.truncate_first
|
|
declared = len(chunk) # honest Content-Length...
|
|
sent = chunk[: declared // 2] if partial else chunk # ...but under-deliver the body
|
|
|
|
if rng:
|
|
self.send_response(206)
|
|
self.send_header("Content-Range", f"bytes {start}-{total - 1}/{total}")
|
|
else:
|
|
self.send_response(200)
|
|
self.send_header("Accept-Ranges", "bytes")
|
|
self.send_header("Content-Length", str(declared))
|
|
self.end_headers()
|
|
self.wfile.write(sent) # short write + return => connection closes (clean FIN)
|
|
|
|
self._httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
|
self._thread = threading.Thread(target=self._httpd.serve_forever, daemon=True)
|
|
|
|
def __enter__(self) -> "_FaultServer":
|
|
self._thread.start()
|
|
return self
|
|
|
|
def __exit__(self, *_exc: object) -> None:
|
|
self._httpd.shutdown()
|
|
self._httpd.server_close()
|
|
|
|
@property
|
|
def url(self) -> str:
|
|
host, port = self._httpd.server_address
|
|
return f"http://{host}:{port}/file"
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def fast_retries(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(dl, "RETRY_WAIT", 0)
|
|
|
|
|
|
def test_truncated_segment_is_not_finalized(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A segment cut short mid-body must raise, never finalize a truncated file."""
|
|
monkeypatch.setattr(dl, "MAX_ATTEMPTS", 2)
|
|
body = b"A" * 2000
|
|
save = tmp_path / "seg.bin"
|
|
with _FaultServer(body) as srv:
|
|
srv.truncate_first = 99 # every attempt under-delivers
|
|
with pytest.raises(Exception):
|
|
consume(dl.download(url=srv.url, save_path=save, session=Session(), segmented=True))
|
|
assert not save.exists(), "truncated segment must not be finalized as complete"
|
|
|
|
|
|
def test_interrupted_download_resumes_from_bytes_written(tmp_path: Path) -> None:
|
|
"""After a mid-body interruption the retry must Range-resume from the real
|
|
byte count (not a pre-allocated full size) and produce the correct file."""
|
|
body = b"".join(bytes([i % 256]) for i in range(4000))
|
|
save = tmp_path / "file.bin"
|
|
with _FaultServer(body) as srv:
|
|
srv.truncate_first = 1 # first attempt delivers half then closes; retry completes
|
|
consume(dl.download(url=srv.url, save_path=save, session=Session()))
|
|
assert save.read_bytes() == body
|
|
# attempt 1 had no Range; attempt 2 must resume partway in, not from 0 or past the end
|
|
assert srv.ranges[0] is None
|
|
resume_start = int(srv.ranges[1].removeprefix("bytes=").split("-", 1)[0])
|
|
assert 0 < resume_start < len(body), f"resume offset {resume_start} not derived from bytes written"
|
|
|
|
|
|
def test_resumes_from_preexisting_dev_partial(tmp_path: Path) -> None:
|
|
"""A .!dev left on disk by a killed process (Ctrl+C, outage, crash) must resume on the
|
|
next run via a Range request from the partial's size, not restart from zero."""
|
|
body = b"".join(bytes([i % 256]) for i in range(4000))
|
|
save = tmp_path / "file.bin"
|
|
tmp = save.with_name("file.bin.!dev")
|
|
tmp.write_bytes(body[:1500]) # correct prefix the interrupted run flushed to disk
|
|
with _FaultServer(body) as srv:
|
|
consume(dl.download(url=srv.url, save_path=save, session=Session()))
|
|
assert save.read_bytes() == body
|
|
assert not tmp.exists()
|
|
assert srv.ranges[0] == "bytes=1500-", f"expected resume from 1500, got {srv.ranges[0]}"
|
|
|
|
|
|
def test_byte_range_slice_segments_never_probed_or_boosted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Adaptive tail boost must skip segments whose per-item headers carry a Range.
|
|
|
|
Such segments are byte-range slices of one larger resource; the boost probe and
|
|
part-mode download would overwrite that Range and pull the entire resource into
|
|
the segment file. A server-side barrier clusters the first six completions so the
|
|
main loop reliably sees spare workers while the two tail slices are still queued,
|
|
the exact window where an unguarded boost would corrupt output.
|
|
"""
|
|
seg = 2 * 1024 * 1024
|
|
nseg = 8
|
|
body = b"".join(bytes([i]) * seg for i in range(nseg)) # 16 MiB, distinct byte per slice
|
|
|
|
barrier_count = [0]
|
|
barrier_lock = threading.Lock()
|
|
barrier = threading.Event()
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def log_message(self, *_a: object) -> None:
|
|
pass
|
|
|
|
def do_GET(self) -> None: # noqa: N802 (stdlib name)
|
|
rng = self.headers.get("Range") or ""
|
|
with barrier_lock:
|
|
barrier_count[0] += 1
|
|
if barrier_count[0] >= 6:
|
|
barrier.set()
|
|
barrier.wait(5) # cluster the first six completions; auto-release as a safety net
|
|
start, _, end_s = rng.removeprefix("bytes=").partition("-")
|
|
start = int(start or 0)
|
|
end = int(end_s) if end_s else len(body) - 1
|
|
chunk = body[start : end + 1]
|
|
self.send_response(206 if rng else 200)
|
|
if rng:
|
|
self.send_header("Content-Range", f"bytes {start}-{end}/{len(body)}")
|
|
self.send_header("Accept-Ranges", "bytes")
|
|
self.send_header("Content-Length", str(len(chunk)))
|
|
self.end_headers()
|
|
self.wfile.write(chunk)
|
|
|
|
httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
|
threading.Thread(target=httpd.serve_forever, daemon=True).start()
|
|
host, port = httpd.server_address
|
|
|
|
probed: list[str] = []
|
|
real_probe = dl.probe_ranged
|
|
monkeypatch.setattr(
|
|
dl, "probe_ranged", lambda url, session, **kw: probed.append(url) or real_probe(url, session, **kw)
|
|
)
|
|
|
|
url = f"http://{host}:{port}/media.mp4"
|
|
urls = [{"url": url, "headers": {"Range": f"bytes={i * seg}-{(i + 1) * seg - 1}"}} for i in range(nseg)]
|
|
try:
|
|
consume(
|
|
dl.requests(
|
|
urls=urls,
|
|
output_dir=tmp_path,
|
|
filename="seg_{i:02}{ext}",
|
|
session=Session(),
|
|
max_workers=16,
|
|
adaptive=True,
|
|
)
|
|
)
|
|
finally:
|
|
httpd.shutdown()
|
|
httpd.server_close()
|
|
|
|
assert probed == [], "byte-range slice segments must never be range-probed"
|
|
for i in range(nseg):
|
|
data = (tmp_path / f"seg_{i:02}.mp4").read_bytes()
|
|
assert data == body[i * seg : (i + 1) * seg], f"segment {i} corrupted: got {len(data)} bytes"
|
|
|
|
|
|
def test_oversized_resume_416_restarts_clean(tmp_path: Path) -> None:
|
|
"""A stale .!dev larger than the resource provokes a 416; the downloader must
|
|
discard it and re-download cleanly rather than burn retries or corrupt."""
|
|
body = b"Z" * 1500
|
|
save = tmp_path / "file.bin"
|
|
tmp = save.with_name(f"{save.name}.!dev")
|
|
tmp.write_bytes(b"X" * 4096) # poisoned partial, bigger than the server's body
|
|
with _FaultServer(body) as srv:
|
|
consume(dl.download(url=srv.url, save_path=save, session=Session()))
|
|
assert save.read_bytes() == body
|
|
assert srv.ranges[0] is not None and "4096" in srv.ranges[0] # first request over-ranged → 416
|