diff --git a/docs/CLI_GUIDE.md b/docs/CLI_GUIDE.md index b984567..238e471 100644 --- a/docs/CLI_GUIDE.md +++ b/docs/CLI_GUIDE.md @@ -66,6 +66,17 @@ reliadl download --url URL --output OUTPUT_PATH [OPTIONS] | `--output`, `-o` | Yes | Target destination file path | | `--adachunk` | No | Enable dynamic AdaChunk Lyapunov chunk size optimization | | `--whittle` | No | Enable Whittle index multi-armed bandit mirror scheduling | +| `--expected-hash`, `--sha256` | No | Expected SHA-256 of the whole file; the download fails if it does not match | +| `--workers`, `-j` | No | Parallel connections, 1–32 (default: `download.max_parallel_workers`, 4) | +| `--chunk-size` | No | Chunk size such as `8MB`, 1MB–256MB (default: `download.chunk_size`) | +| `--limit-rate` | No | Bandwidth cap per second such as `10MB`; `0` is unlimited | +| `--config` | No | Path to a ReliaDL YAML configuration file | + +**How it works**: the server is probed with `HEAD` (or a one-byte range `GET` if `HEAD` is refused) for its size, `Accept-Ranges`, and `ETag`. The file is split into chunks, and a pool of workers fetches them concurrently with `Range` requests, retrying transient failures (timeouts, resets, 429, 5xx) with exponential backoff. Bytes are written straight into a pre-allocated `.part` file and each chunk is SHA-256 hashed as it streams. After every chunk completes, the session is checkpointed to `.ReliaDL/.state` next to the output. When all chunks are done the whole file is hashed and `.part` is renamed to the output path. + +Servers that do not support byte ranges are rejected for now; single-stream fallback is tracked in #42. + +**Interrupting**: `Ctrl+C` (or `SIGTERM`) stops the transfer, saves the checkpoint, prints the `resume` command, and exits with status `130`. **Example**: ```bash @@ -92,10 +103,13 @@ reliadl resume --state-file STATE_FILE_PATH | Parameter | Required | Description | |---|---|---| | `--state-file` | Yes | Absolute or relative path to `.state` checkpoint file | +| `--workers`, `--chunk-size`, `--limit-rate`, `--config` | No | As for `download`; the chunk layout recorded in the state file is kept | + +Before continuing, the remote file is probed again: if its size or `ETag` has changed, the resume is refused. Every chunk the state file marks complete is re-hashed from `.part`, and any that no longer match are downloaded again. **Example**: ```bash -reliadl resume --state-file "./downloads/.reliadl/models_v2.tar.gz.state" +reliadl resume --state-file "./downloads/.ReliaDL/models_v2.tar.gz.state" ``` --- diff --git a/pyproject.toml b/pyproject.toml index 582e32b..290c949 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,9 +28,14 @@ dependencies = [ "PyYAML>=6.0.3", "structlog>=26.1.0", "jsonschema>=4.26.0", - "cryptography>=50.0.1" + "cryptography>=50.0.1", + "httpx>=0.28.1" ] +[project.optional-dependencies] +http2 = ["httpx[http2]>=0.28.1"] +socks = ["httpx[socks]>=0.28.1"] + [project.scripts] reliadl = "reliadl.cli:main" diff --git a/reliadl/__init__.py b/reliadl/__init__.py index 58bca19..e432e7b 100644 --- a/reliadl/__init__.py +++ b/reliadl/__init__.py @@ -56,6 +56,7 @@ ConfigurationError, ConnectionError, DiskFullError, + DownloadCancelledError, FileHashMismatchError, HTTPError, IntegrityError, @@ -68,6 +69,7 @@ ProxyAuthenticationError, ProxyConnectionError, ProxyError, + RangeNotSupportedError, ReliaDLError, ServerError, StateCorruptedError, @@ -78,6 +80,11 @@ SubBlockCorruptedError, TimeoutError, ) +from reliadl.download_engine import ( + DownloadEngine, + RemoteFileInfo, + plan_chunks, +) from reliadl.file_assembler import ( DEFAULT_ASSEMBLY_BUFFER_SIZE, FileAssembler, @@ -189,6 +196,12 @@ "StateCorruptedError", "AssemblyError", "AssemblyFailedError", + "RangeNotSupportedError", + "DownloadCancelledError", + # Download Engine + "DownloadEngine", + "RemoteFileInfo", + "plan_chunks", # Config "parse_size", "format_size", diff --git a/reliadl/cli.py b/reliadl/cli.py index 39bb125..e3b9bf4 100644 --- a/reliadl/cli.py +++ b/reliadl/cli.py @@ -8,21 +8,26 @@ from __future__ import annotations import argparse +import asyncio import json import os import platform +import signal import sys import time import urllib.request import urllib.error from pathlib import Path -from typing import Optional, Sequence +from typing import Any, Awaitable, Callable, Optional, Sequence from reliadl.adapters.proxy_adapter import ProxyConfig, ProxyTunnel -from reliadl.config import format_size +from reliadl.config import format_size, get_download_config +from reliadl.download_engine import DownloadEngine +from reliadl.exceptions import DownloadCancelledError, ReliaDLError from reliadl.hash_verifier import StreamingHashVerifier, compute_file_hash, constant_time_compare from reliadl.logger import configure_logger, get_logger from reliadl.manifest import BinaryMerkleTree, compute_merkle_root, load_manifest +from reliadl.models import DownloadConfig, DownloadResult, ProgressReport from reliadl.state_manager import StateManager logger = get_logger("reliadl.cli") @@ -478,6 +483,85 @@ def run_hash_tree(args: argparse.Namespace) -> int: # Core Commands: Download, Resume, Verify # ───────────────────────────────────────────────────────────────────────────── +def _engine_config(args: argparse.Namespace) -> DownloadConfig: + """Load the layered YAML/env configuration with command-line overrides on top.""" + download: dict[str, Any] = {} + network: dict[str, Any] = {} + if getattr(args, "workers", None) is not None: + download["max_parallel_workers"] = args.workers + if getattr(args, "chunk_size", None) is not None: + download["chunk_size"] = args.chunk_size + if getattr(args, "limit_rate", None) is not None: + network["max_bandwidth"] = args.limit_rate + return get_download_config( + config_path=getattr(args, "config", None), + overrides={"download": download, "network": network}, + ) + + +def _print_progress(report: ProgressReport) -> None: + """Render a single self-overwriting progress line on an interactive stderr.""" + eta = report.estimated_remaining_seconds + line = ( + f"\r{report.percentage:6.2f}% " + f"{format_size(report.downloaded_bytes)} / {format_size(report.total_bytes)} " + f"{format_size(int(report.current_speed_bps))}/s " + f"chunks {report.chunks_complete}/{report.total_chunks} " + f"ETA {f'{eta:.0f}s' if eta is not None else '--'} " + ) + sys.stderr.write(line) + sys.stderr.flush() + + +def _run_engine_session( + engine: DownloadEngine, + session: Callable[[], Awaitable[DownloadResult]], +) -> int: + """Run a download session, turning SIGINT/SIGTERM into a checkpointed stop.""" + + async def runner() -> DownloadResult: + loop = asyncio.get_running_loop() + installed = [] + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, engine.request_shutdown) + installed.append(sig) + except (NotImplementedError, RuntimeError, ValueError): + # Windows event loops have no signal handlers; Ctrl+C there + # cancels the session task, which checkpoints the same way. + pass + try: + return await session() + finally: + for sig in installed: + loop.remove_signal_handler(sig) + + try: + result = asyncio.run(runner()) + except DownloadCancelledError as err: + print(f"\n[INFO] {err.message}") + return 130 + except KeyboardInterrupt: + print("\n[INFO] Download interrupted; state saved for 'reliadl resume'.") + return 130 + except ReliaDLError as err: + print(f"\n[ERROR] {err.message}") + return 1 + finally: + if sys.stderr.isatty(): + sys.stderr.write("\n") + + print(f"[SUCCESS] Saved {format_size(result.file_size)} to: {result.output_path}") + if result.file_hash: + status = "verified" if result.is_verified else "computed" + print(f"SHA-256 ({status}): {result.file_hash}") + print( + f"[INFO] {result.total_chunks} chunks in {result.elapsed_seconds:.2f}s " + f"({format_size(int(result.average_speed_bps))}/s, {result.chunks_retried} retried)" + ) + return 0 + + def run_download(args: argparse.Namespace) -> int: """Execute high-speed chunked parallel file download.""" url = args.url @@ -485,14 +569,14 @@ def run_download(args: argparse.Namespace) -> int: print(f"[INFO] Initiating parallel download from: {url}") print(f"[INFO] Target output path: {output}") - output.parent.mkdir(parents=True, exist_ok=True) - - # Perform lightweight download initialization - state_mgr = StateManager(output) - state_mgr.initialize(file_size_bytes=1024 * 1024, chunk_size_bytes=256 * 1024) + try: + config = _engine_config(args) + except (ReliaDLError, ValueError) as err: + print(f"[ERROR] Invalid configuration: {err}") + return 1 - print(f"[SUCCESS] Download target initialized cleanly at: {output}") - return 0 + engine = DownloadEngine(config, progress_callback=_print_progress if sys.stderr.isatty() else None) + return _run_engine_session(engine, lambda: engine.download(url, output, expected_hash=args.expected_hash)) def run_resume(args: argparse.Namespace) -> int: @@ -502,7 +586,15 @@ def run_resume(args: argparse.Namespace) -> int: print(f"[ERROR] State file not found: {state_file}") return 1 print(f"[INFO] Resuming transfer session from state file: {state_file}") - return 0 + + try: + config = _engine_config(args) + except (ReliaDLError, ValueError) as err: + print(f"[ERROR] Invalid configuration: {err}") + return 1 + + engine = DownloadEngine(config, progress_callback=_print_progress if sys.stderr.isatty() else None) + return _run_engine_session(engine, lambda: engine.resume(state_file)) def run_verify(args: argparse.Namespace) -> int: @@ -531,6 +623,14 @@ def run_verify(args: argparse.Namespace) -> int: # CLI Parser Definition & Main Dispatcher # ───────────────────────────────────────────────────────────────────────────── +def _add_engine_arguments(parser: argparse.ArgumentParser) -> None: + """Options shared by the commands that drive the download engine.""" + parser.add_argument("--workers", "-j", type=int, help="Parallel connections (1-32, default from config: 4)") + parser.add_argument("--chunk-size", help="Chunk size, e.g. 8MB (1MB-256MB)") + parser.add_argument("--limit-rate", help="Bandwidth cap per second, e.g. 10MB (0 = unlimited)") + parser.add_argument("--config", help="Path to a ReliaDL YAML configuration file") + + def build_parser() -> argparse.ArgumentParser: """Construct argument parser for ReliaDL CLI.""" parser = argparse.ArgumentParser( @@ -548,10 +648,13 @@ def build_parser() -> argparse.ArgumentParser: p_dl.add_argument("--output", "-o", required=True, help="Output file path") p_dl.add_argument("--adachunk", action="store_true", help="Enable AdaChunk dynamic chunk optimization") p_dl.add_argument("--whittle", action="store_true", help="Enable Whittle index mirror bandit scheduling") + p_dl.add_argument("--expected-hash", "--sha256", dest="expected_hash", help="Expected SHA-256 of the whole file") + _add_engine_arguments(p_dl) # Core: resume p_res = subparsers.add_parser("resume", help="Resume interrupted transfer session") p_res.add_argument("--state-file", required=True, help="Path to .state file") + _add_engine_arguments(p_res) # Core: verify p_ver = subparsers.add_parser("verify", help="Verify payload SHA-256 hash digest") diff --git a/reliadl/download_engine.py b/reliadl/download_engine.py new file mode 100644 index 0000000..5e63d65 --- /dev/null +++ b/reliadl/download_engine.py @@ -0,0 +1,714 @@ +""" +Asynchronous HTTP range download engine for ReliaDL. + +Probes the origin for range support, splits the file into fixed-size chunks, +and fetches them concurrently over a pooled ``httpx.AsyncClient``. Each chunk +is streamed straight into a pre-allocated ``.part`` file with +positional writes, hashed on the fly, and checkpointed to the session state +file once its bytes are on disk. An interrupted session (Ctrl+C, SIGTERM, a +chunk that exhausts its retries) leaves a state file that ``resume`` picks up, +re-verifying every chunk it claims to have before trusting it. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import os +import random +import time +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable, Optional, Union + +import httpx + +from reliadl.exceptions import ( + ChunkHashMismatchError, + ClientError, + ConnectionError, + DownloadCancelledError, + FileHashMismatchError, + HTTPError, + NetworkError, + PreconditionFailedError, + RangeNotSupportedError, + ReliaDLError, + ServerError, + TimeoutError, +) +from reliadl.hash_verifier import ( + StreamingHashVerifier, + compute_file_hash, + constant_time_compare, + normalize_hash, +) +from reliadl.logger import get_logger +from reliadl.models import ( + ChunkState, + ChunkStatus, + DownloadConfig, + DownloadResult, + DownloadState, + DownloadStatus, + ProgressReport, +) +from reliadl.rate_limiter import TokenBucketRateLimiter +from reliadl.sparse_writer import SparseFileWriter +from reliadl.state_manager import StateManager + +logger = get_logger("reliadl.download_engine") + +PART_SUFFIX = ".part" +_REVALIDATE_READ_SIZE = 1024 * 1024 + +ProgressCallback = Callable[[ProgressReport], None] + + +@dataclass(frozen=True) +class RemoteFileInfo: + """What a HEAD (or one-byte range) probe learned about the remote file.""" + + url: str + size: Optional[int] + accepts_ranges: bool + etag: Optional[str] = None + last_modified: Optional[str] = None + + @property + def if_range(self) -> Optional[str]: + """Validator for ``If-Range``; RFC 9110 only allows a strong ETag or a date.""" + if self.etag and not self.etag.startswith("W/"): + return self.etag + return self.last_modified + + +def plan_chunks(file_size: int, chunk_size: int) -> list[ChunkState]: + """Split ``file_size`` bytes into contiguous, non-overlapping chunk states.""" + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + return [ + ChunkState(index=index, start_byte=start, end_byte=min(start + chunk_size, file_size) - 1) + for index, start in enumerate(range(0, file_size, chunk_size)) + ] + + +def _parse_int(value: Optional[str]) -> Optional[int]: + try: + return int(value) if value is not None else None + except ValueError: + return None + + +def _parse_content_range(value: Optional[str]) -> Optional[tuple[int, int, Optional[int]]]: + """Parse ``bytes -/`` into ``(start, end, total)``.""" + if not value or not value.lower().startswith("bytes "): + return None + try: + span, _, total = value[6:].strip().partition("/") + start, _, end = span.partition("-") + return int(start), int(end), None if total in ("", "*") else int(total) + except ValueError: + return None + + +def _retry_after_seconds(response: httpx.Response) -> Optional[float]: + value = response.headers.get("retry-after") + if value is None: + return None + try: + return max(0.0, float(value)) + except ValueError: + return None + + +def _http_error(response: httpx.Response, url: str) -> HTTPError: + """Map a non-success response onto the ReliaDL HTTP error hierarchy.""" + status = response.status_code + message = f"HTTP {status} {response.reason_phrase} from {url}" + headers = dict(response.headers) + if status == 412: + return PreconditionFailedError(message, etag=response.headers.get("etag"), url=url) + if 400 <= status < 500: + return ClientError( + message, + retry_after=_retry_after_seconds(response), + status_code=status, + response_headers=headers, + url=url, + ) + if status >= 500: + return ServerError(message, status_code=status, response_headers=headers, url=url) + return HTTPError(message, status_code=status, response_headers=headers, url=url, is_retryable=False) + + +def _transport_error(err: httpx.HTTPError, url: str) -> NetworkError: + """Map an httpx transport failure onto the ReliaDL network error hierarchy.""" + if isinstance(err, httpx.TimeoutException): + return TimeoutError(f"Timed out talking to {url}: {err!r}", url=url, cause=err) + if isinstance(err, httpx.ConnectError): + return ConnectionError(f"Could not connect to {url}: {err}", host=httpx.URL(url).host, url=url, cause=err) + return NetworkError(f"Transfer from {url} failed: {err!r}", url=url, cause=err) + + +class DownloadEngine: + """ + Parallel range downloader bounded by ``config.max_parallel_workers``. + + One engine runs one session at a time. ``request_shutdown`` may be called + from a signal handler on the engine's event loop to stop the session + cleanly; the awaiting ``download``/``resume`` call then raises + ``DownloadCancelledError`` after the state file has been written. + """ + + def __init__( + self, + config: Optional[DownloadConfig] = None, + *, + state_manager: Optional[StateManager] = None, + progress_callback: Optional[ProgressCallback] = None, + transport: Optional[httpx.AsyncBaseTransport] = None, + ) -> None: + self._config = config if config is not None else DownloadConfig() + self._state_manager = state_manager if state_manager is not None else StateManager( + default_state_dir=self._config.state_directory, + ) + self._progress_callback = progress_callback + self._transport = transport + self._shutdown = asyncio.Event() + self._limiter: Optional[TokenBucketRateLimiter] = None + if self._config.max_bandwidth_bytes_per_sec > 0: + self._limiter = TokenBucketRateLimiter(float(self._config.max_bandwidth_bytes_per_sec)) + # Per-session progress accounting; bytes of chunks still in flight are + # tracked separately so a retried chunk does not count twice. + self._in_flight: dict[int, int] = {} + self._state: Optional[DownloadState] = None + self._reset_progress() + + @property + def config(self) -> DownloadConfig: + return self._config + + def request_shutdown(self) -> None: + """Stop dispatching chunks, abort in-flight ones, and checkpoint state.""" + self._shutdown.set() + + # ── Public entry points ──────────────────────────────────────────────── + + async def download( + self, + url: str, + output_path: Union[str, Path], + expected_hash: Optional[str] = None, + ) -> DownloadResult: + """Download ``url`` to ``output_path``, starting a fresh session.""" + output = Path(output_path).resolve() + part_path = output.with_name(output.name + PART_SUFFIX) + expected = normalize_hash(expected_hash) if expected_hash else None + + async with self._build_client() as client: + info = await self.probe(client, url) + if not info.accepts_ranges or info.size is None: + raise RangeNotSupportedError( + f"{url} does not support byte-range requests " + "(no 'Accept-Ranges: bytes' and no 206 reply to a range probe)", + url=url, + ) + + state = DownloadState( + download_id=uuid.uuid4().hex, + url=url, + target_path=str(output), + file_size=info.size, + chunk_size=self._config.chunk_size_bytes, + hash_algorithm=self._config.hash_algorithm, + expected_file_hash=expected, + etag=info.etag, + output_path=str(part_path), + status=DownloadStatus.DOWNLOADING, + chunks=plan_chunks(info.size, self._config.chunk_size_bytes), + ) + # A leftover .part from an abandoned session holds unrelated bytes. + part_path.unlink(missing_ok=True) + writer = SparseFileWriter( + part_path, + info.size, + check_disk_space=self._config.pre_check_disk_space, + ) + state_path = self._state_manager.get_state_path(output) + return await self._run(client, info, state, state_path, writer) + + async def resume(self, state_path: Union[str, Path]) -> DownloadResult: + """Continue the session recorded in ``state_path``.""" + path = Path(state_path).resolve() + state = self._state_manager.load(path) + part_path = Path(state.output_path or state.target_path + PART_SUFFIX) + + async with self._build_client() as client: + info = await self.probe(client, state.url) + if not info.accepts_ranges or info.size is None: + raise RangeNotSupportedError( + f"{state.url} no longer supports byte-range requests", + url=state.url, + ) + if info.size != state.file_size or ( + state.etag and info.etag and state.etag != info.etag + ): + raise PreconditionFailedError( + f"Remote file changed since the session started " + f"(size {state.file_size} -> {info.size}, " + f"ETag {state.etag} -> {info.etag}); start a new download", + etag=info.etag, + url=state.url, + ) + + if not part_path.is_file(): + for chunk in state.chunks: + self._reset_chunk(chunk) + writer = SparseFileWriter(part_path, state.file_size, check_disk_space=False) + await asyncio.to_thread(self._revalidate_chunks, state, writer) + state.status = DownloadStatus.DOWNLOADING + return await self._run(client, info, state, path, writer) + + async def probe(self, client: httpx.AsyncClient, url: str) -> RemoteFileInfo: + """ + Discover size, range support, and validators for ``url``. + + HEAD is tried first. Servers that reject HEAD or do not advertise + ``Accept-Ranges`` get a one-byte range GET, whose 206 and + ``Content-Range`` total settle both questions at once. + """ + try: + head = await client.head(url) + except httpx.HTTPError as err: + raise _transport_error(err, url) from err + + if head.is_success: + size = _parse_int(head.headers.get("content-length")) + if size is not None and head.headers.get("accept-ranges", "").lower() == "bytes": + return RemoteFileInfo( + url=str(head.url), + size=size, + accepts_ranges=True, + etag=head.headers.get("etag"), + last_modified=head.headers.get("last-modified"), + ) + + try: + async with client.stream("GET", url, headers={"Range": "bytes=0-0"}) as response: + etag = response.headers.get("etag") + last_modified = response.headers.get("last-modified") + if response.status_code == 206: + parsed = _parse_content_range(response.headers.get("content-range")) + total = parsed[2] if parsed else None + return RemoteFileInfo( + url=str(response.url), + size=total, + accepts_ranges=total is not None, + etag=etag, + last_modified=last_modified, + ) + if not response.is_success: + raise _http_error(response, url) + return RemoteFileInfo( + url=str(response.url), + size=_parse_int(response.headers.get("content-length")), + accepts_ranges=False, + etag=etag, + last_modified=last_modified, + ) + except httpx.HTTPError as err: + raise _transport_error(err, url) from err + + # ── Session orchestration ────────────────────────────────────────────── + + async def _run( + self, + client: httpx.AsyncClient, + info: RemoteFileInfo, + state: DownloadState, + state_path: Path, + writer: SparseFileWriter, + ) -> DownloadResult: + self._shutdown.clear() + self._reset_progress() + self._total_bytes = state.file_size + self._completed_bytes = sum(c.size for c in state.completed_chunks) + self._resumed_bytes = self._completed_bytes + self._state = state + + queue: asyncio.Queue[ChunkState] = asyncio.Queue() + for chunk in state.chunks: + if not chunk.status.is_successful: + self._reset_chunk(chunk) + queue.put_nowait(chunk) + + self._state_manager.save(state, state_path) + worker_count = min(self._config.max_parallel_workers, max(1, queue.qsize())) + logger.info( + "download.start", + url=state.url, + size=state.file_size, + chunks=state.total_chunks, + pending=queue.qsize(), + workers=worker_count, + ) + + workers = [ + asyncio.create_task(self._worker(client, info, queue, state, state_path, writer)) + for _ in range(worker_count if queue.qsize() else 0) + ] + shutdown_waiter = asyncio.create_task(self._shutdown.wait()) + try: + running = set(workers) + while running: + await asyncio.wait({*running, shutdown_waiter}, return_when=asyncio.FIRST_COMPLETED) + running = {w for w in workers if not w.done()} + failed = next((w for w in workers if w.done() and not w.cancelled() and w.exception()), None) + if failed is not None: + await self._stop_workers(workers) + self._checkpoint_interrupted(state, state_path, writer, status=DownloadStatus.FAILED) + raise failed.exception() # type: ignore[misc] + if running and self._shutdown.is_set(): + await self._stop_workers(workers) + self._checkpoint_interrupted(state, state_path, writer) + raise DownloadCancelledError( + f"Download interrupted; resume with: reliadl resume --state-file {state_path}", + state_file=str(state_path), + ) + except asyncio.CancelledError: + await self._stop_workers(workers) + self._checkpoint_interrupted(state, state_path, writer) + raise + finally: + shutdown_waiter.cancel() + + return await self._finalize(state, state_path, writer) + + async def _worker( + self, + client: httpx.AsyncClient, + info: RemoteFileInfo, + queue: asyncio.Queue[ChunkState], + state: DownloadState, + state_path: Path, + writer: SparseFileWriter, + ) -> None: + while not self._shutdown.is_set(): + try: + chunk = queue.get_nowait() + except asyncio.QueueEmpty: + return + self._active_workers += 1 + try: + await self._download_chunk_with_retries(client, info, chunk, state.hash_algorithm, writer) + finally: + self._active_workers -= 1 + # The state file may only claim bytes that are durable on disk. + writer.sync() + self._state_manager.save(state, state_path) + + async def _download_chunk_with_retries( + self, + client: httpx.AsyncClient, + info: RemoteFileInfo, + chunk: ChunkState, + algorithm: str, + writer: SparseFileWriter, + ) -> None: + max_attempts = max(1, self._config.max_retries_per_chunk) + while True: + chunk.mark_in_progress() + try: + digest = await self._fetch_chunk(client, info, chunk, algorithm, writer) + except ReliaDLError as err: + self._in_flight.pop(chunk.index, None) + if not err.is_retryable or chunk.retries + 1 >= max_attempts: + chunk.mark_abandoned(str(err)) + logger.error("chunk.abandoned", chunk=chunk.index, attempts=chunk.retries + 1, error=str(err)) + raise + chunk.mark_failed(str(err)) + delay = self._backoff_delay(chunk.retries, err) + logger.warning("chunk.retry", chunk=chunk.index, attempt=chunk.retries, delay=round(delay, 2), error=str(err)) + await asyncio.sleep(delay) + continue + chunk.mark_complete(digest) + self._in_flight.pop(chunk.index, None) + self._completed_bytes += chunk.size + self._emit_progress(force=True) + return + + async def _fetch_chunk( + self, + client: httpx.AsyncClient, + info: RemoteFileInfo, + chunk: ChunkState, + algorithm: str, + writer: SparseFileWriter, + ) -> str: + headers = {"Range": f"bytes={chunk.start_byte}-{chunk.end_byte}"} + if info.if_range: + headers["If-Range"] = info.if_range + + hasher = StreamingHashVerifier(algorithm) + offset = chunk.start_byte + self._in_flight[chunk.index] = 0 + try: + async with client.stream("GET", info.url, headers=headers) as response: + if response.status_code == 200: + # A 200 to a ranged request means the server ignored the + # range, or If-Range failed because the file changed. + if info.if_range: + raise PreconditionFailedError( + f"Remote file changed during download (If-Range {info.if_range} no longer matches)", + etag=response.headers.get("etag"), + url=info.url, + ) + raise RangeNotSupportedError( + f"Server ignored the byte range for chunk {chunk.index}", + url=info.url, + ) + if response.status_code != 206: + raise _http_error(response, info.url) + + parsed = _parse_content_range(response.headers.get("content-range")) + if parsed is None or parsed[:2] != (chunk.start_byte, chunk.end_byte): + raise HTTPError( + f"Chunk {chunk.index}: requested bytes {chunk.start_byte}-{chunk.end_byte}, " + f"server sent Content-Range {response.headers.get('content-range')!r}", + status_code=206, + url=info.url, + is_retryable=False, + ) + + async for data in response.aiter_raw(): + if offset + len(data) > chunk.end_byte + 1: + raise HTTPError( + f"Chunk {chunk.index}: server sent more bytes than the requested range", + status_code=206, + url=info.url, + is_retryable=False, + ) + await self._throttle(len(data)) + writer.write_at(offset, data) + hasher.update(data) + offset += len(data) + self._in_flight[chunk.index] = offset - chunk.start_byte + self._emit_progress() + except httpx.HTTPError as err: + raise _transport_error(err, info.url) from err + + received = offset - chunk.start_byte + if received != chunk.size: + raise NetworkError( + f"Chunk {chunk.index} truncated: received {received} of {chunk.size} bytes", + url=info.url, + ) + + digest = hasher.hexdigest() + if chunk.expected_hash and not constant_time_compare(digest, chunk.expected_hash): + raise ChunkHashMismatchError( + f"Chunk {chunk.index} failed hash verification", + chunk_index=chunk.index, + expected_hash=chunk.expected_hash, + computed_hash=digest, + start_byte=chunk.start_byte, + end_byte=chunk.end_byte, + ) + return digest + + async def _finalize( + self, + state: DownloadState, + state_path: Path, + writer: SparseFileWriter, + ) -> DownloadResult: + await asyncio.to_thread(writer.sync) + writer.close() + part_path = writer.target_path + output = Path(state.target_path) + + file_hash = "" + verified = False + if self._config.verify_on_complete or state.expected_file_hash: + state.status = DownloadStatus.VERIFYING + file_hash = await asyncio.to_thread(compute_file_hash, part_path, state.hash_algorithm) + if state.expected_file_hash: + if not constant_time_compare(file_hash, state.expected_file_hash): + state.status = DownloadStatus.FAILED + self._state_manager.save(state, state_path) + raise FileHashMismatchError( + f"Downloaded file hash {file_hash} does not match expected {state.expected_file_hash}", + context={"path": str(part_path)}, + ) + verified = True + + os.replace(part_path, output) + self._state_manager.delete(state_path) + try: + state_path.parent.rmdir() # only succeeds when no other session uses it + except OSError as exc: + logger.debug("download.cleanup_state_dir_failed", path=str(state_path.parent), error=str(exc)) + state.status = DownloadStatus.COMPLETE + self._emit_progress(force=True) + + elapsed = time.monotonic() - self._started_at + fetched = state.file_size - self._resumed_bytes + logger.info("download.complete", path=str(output), size=state.file_size, seconds=round(elapsed, 3)) + return DownloadResult( + output_path=output, + file_hash=file_hash, + is_verified=verified, + file_size=state.file_size, + total_chunks=state.total_chunks, + chunks_retried=sum(1 for c in state.chunks if c.retries), + total_bytes_downloaded=fetched, + elapsed_seconds=elapsed, + average_speed_bps=fetched / elapsed if elapsed > 0 else 0.0, + download_id=state.download_id, + ) + + # ── Helpers ──────────────────────────────────────────────────────────── + + def _build_client(self) -> httpx.AsyncClient: + cfg = self._config + return httpx.AsyncClient( + transport=self._transport, + http2=cfg.http2 and importlib.util.find_spec("h2") is not None, + verify=cfg.verify_ssl, + proxy=cfg.proxy_url, + follow_redirects=True, + max_redirects=cfg.max_redirects, + timeout=httpx.Timeout( + connect=cfg.connect_timeout_seconds, + read=cfg.read_timeout_seconds, + write=cfg.read_timeout_seconds, + pool=None, + ), + limits=httpx.Limits( + max_connections=cfg.max_parallel_workers, + max_keepalive_connections=cfg.max_parallel_workers, + ), + # Content-Length and byte offsets must describe the stored bytes, + # not a compressed encoding of them. + headers={"User-Agent": cfg.user_agent, "Accept-Encoding": "identity"}, + ) + + async def _stop_workers(self, workers: list[asyncio.Task[None]]) -> None: + for worker in workers: + worker.cancel() + await asyncio.gather(*workers, return_exceptions=True) + + def _checkpoint_interrupted( + self, + state: DownloadState, + state_path: Path, + writer: SparseFileWriter, + status: DownloadStatus = DownloadStatus.CANCELLED, + ) -> None: + for chunk in state.chunks: + if chunk.status in (ChunkStatus.IN_PROGRESS, ChunkStatus.DOWNLOADING, ChunkStatus.FAILED): + chunk.status = ChunkStatus.PENDING + state.status = status + writer.sync() + writer.close() + self._state_manager.save(state, state_path) + logger.info("download.checkpoint", state_file=str(state_path), status=status.value) + + def _revalidate_chunks(self, state: DownloadState, writer: SparseFileWriter) -> None: + """Re-hash chunks the state file calls complete; only matching ones are kept.""" + for chunk in state.chunks: + if not chunk.status.is_successful: + continue + if not chunk.computed_hash: + self._reset_chunk(chunk) + continue + hasher = StreamingHashVerifier(state.hash_algorithm) + for offset in range(chunk.start_byte, chunk.end_byte + 1, _REVALIDATE_READ_SIZE): + hasher.update(writer.read_at(offset, min(_REVALIDATE_READ_SIZE, chunk.end_byte + 1 - offset))) + if not hasher.verify(chunk.computed_hash): + logger.warning("chunk.revalidation_failed", chunk=chunk.index) + self._reset_chunk(chunk) + + @staticmethod + def _reset_chunk(chunk: ChunkState) -> None: + if chunk.status == ChunkStatus.ABANDONED: + chunk.retries = 0 + chunk.status = ChunkStatus.PENDING + chunk.hash_verified = False + chunk.computed_hash = None + + def _backoff_delay(self, attempt: int, err: ReliaDLError) -> float: + cfg = self._config + delay = min( + cfg.retry_max_delay_seconds, + cfg.retry_base_delay_seconds * cfg.retry_backoff_factor ** max(0, attempt - 1), + ) + delay += random.uniform(0.0, cfg.retry_jitter_factor * delay) + retry_after = getattr(err, "retry_after", None) + if retry_after is not None: + delay = max(delay, min(float(retry_after), cfg.retry_max_delay_seconds)) + return delay + + async def _throttle(self, amount: int) -> None: + if self._limiter is None: + return + # A single read can exceed the bucket when the rate is very low. + step = max(1, int(self._limiter.capacity_bytes)) + while amount > 0: + take = min(step, amount) + await self._limiter.acquire(take) + amount -= take + + def _reset_progress(self) -> None: + self._started_at = time.monotonic() + self._total_bytes = 0 + self._completed_bytes = 0 + self._resumed_bytes = 0 + self._in_flight = {} + self._active_workers = 0 + self._last_emit = 0.0 + self._last_emit_bytes = 0 + self._state = None + + def _emit_progress(self, force: bool = False) -> None: + if self._progress_callback is None or self._state is None: + return + now = time.monotonic() + if not force and now - self._last_emit < self._config.progress_update_interval_seconds: + return + + downloaded = self._completed_bytes + sum(self._in_flight.values()) + elapsed = now - self._started_at + interval = now - self._last_emit if self._last_emit else elapsed + fetched = downloaded - self._resumed_bytes + current_speed = (downloaded - self._last_emit_bytes) / interval if self._last_emit and interval > 0 else ( + fetched / elapsed if elapsed > 0 else 0.0 + ) + average_speed = fetched / elapsed if elapsed > 0 else 0.0 + remaining = self._total_bytes - downloaded + self._last_emit = now + self._last_emit_bytes = downloaded + + chunks = self._state.chunks + self._progress_callback( + ProgressReport( + download_id=self._state.download_id, + timestamp=datetime.now(timezone.utc), + total_bytes=self._total_bytes, + downloaded_bytes=downloaded, + percentage=100.0 * downloaded / self._total_bytes if self._total_bytes else 100.0, + total_chunks=len(chunks), + chunks_complete=sum(1 for c in chunks if c.status.is_successful), + chunks_in_progress=sum(1 for c in chunks if c.status == ChunkStatus.IN_PROGRESS), + chunks_failed=sum(1 for c in chunks if c.status in (ChunkStatus.FAILED, ChunkStatus.ABANDONED)), + chunks_pending=sum(1 for c in chunks if c.status == ChunkStatus.PENDING), + current_speed_bps=current_speed, + average_speed_bps=average_speed, + elapsed_seconds=elapsed, + estimated_remaining_seconds=remaining / average_speed if average_speed > 0 else None, + active_workers=self._active_workers, + ) + ) diff --git a/reliadl/exceptions.py b/reliadl/exceptions.py index c0a79fb..d589754 100644 --- a/reliadl/exceptions.py +++ b/reliadl/exceptions.py @@ -581,3 +581,32 @@ def __init__( self.reason = reason self.missing_chunks = missing_chunks or [] super().__init__(message, context=ctx, **kwargs) + + +# ───────────────────────────────────────────────────────────────────────────── +# Download Session Hierarchy +# ───────────────────────────────────────────────────────────────────────────── + +class RangeNotSupportedError(NetworkError): + """Raised when the origin cannot serve byte ranges or does not report a file size.""" + + default_retryable: bool = False + + +class DownloadCancelledError(ReliaDLError): + """Raised when a download is interrupted on request after its state has been saved.""" + + default_retryable: bool = False + + def __init__( + self, + message: str, + state_file: Optional[str] = None, + context: Optional[dict[str, Any]] = None, + **kwargs: Any, + ) -> None: + ctx = context.copy() if context else {} + if state_file is not None: + ctx["state_file"] = state_file + self.state_file = state_file + super().__init__(message, context=ctx, **kwargs) diff --git a/requirements.txt b/requirements.txt index f76ba9a..856d6fd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ PyYAML>=6.0.3 structlog>=26.1.0 jsonschema>=4.26.0 cryptography>=50.0.1 +httpx>=0.28.1 diff --git a/tests/unit/test_download_engine.py b/tests/unit/test_download_engine.py new file mode 100644 index 0000000..a65b81e --- /dev/null +++ b/tests/unit/test_download_engine.py @@ -0,0 +1,381 @@ +""" +Unit tests for the asynchronous range download engine in reliadl.download_engine. + +A real threaded HTTP server on localhost serves byte ranges, so the tests cover +the httpx client, the probe, the worker pool, positional writes, checkpointing, +graceful shutdown, and resume end to end. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import io +import os +import tempfile +import threading +import time +import unittest +from contextlib import redirect_stdout +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Optional + +from reliadl.cli import main as cli_main +from reliadl.download_engine import ( + PART_SUFFIX, + DownloadEngine, + _parse_content_range, + plan_chunks, +) +from reliadl.exceptions import ( + ClientError, + DownloadCancelledError, + FileHashMismatchError, + PreconditionFailedError, + RangeNotSupportedError, +) +from reliadl.models import ChunkStatus, DownloadConfig, DownloadStatus, ProgressReport +from reliadl.state_manager import StateManager + +MB = 1024 * 1024 + + +class _RangeServer(ThreadingHTTPServer): + daemon_threads = True + request_queue_size = 64 + + def __init__(self, payload: bytes) -> None: + super().__init__(("127.0.0.1", 0), _RangeHandler) + self.payload = payload + self.etag = '"v1"' + self.allow_head = True + self.ranges = True + self.chunk_delay = 0.0 + # start byte -> HTTP status to answer once, then serve normally + self.fail_once: dict[int, int] = {} + # ranged GETs at or beyond this offset wait until the gate opens + self.gate_from: Optional[int] = None + self.gate = threading.Event() + self.ranged_bytes_served = 0 + self._lock = threading.Lock() + self._active = 0 + self.max_active = 0 + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}/file.bin" + + def handle_error(self, request, client_address) -> None: # type: ignore[no-untyped-def] + # Cancelled workers drop their connections mid-body; that is expected. + pass + + +class _RangeHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + server: _RangeServer + + def log_message(self, format: str, *args: object) -> None: + pass + + def do_HEAD(self) -> None: + if not self.server.allow_head: + self._send(405, b"", head=True) + return + self._serve(head=True) + + def do_GET(self) -> None: + self._serve(head=False) + + def _send(self, status: int, body: bytes, head: bool, headers: Optional[dict[str, str]] = None) -> None: + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + for key, value in (headers or {}).items(): + self.send_header(key, value) + self.end_headers() + if not head: + self.wfile.write(body) + + def _serve(self, head: bool) -> None: + srv = self.server + payload = srv.payload + headers = {"ETag": srv.etag} + if srv.ranges: + headers["Accept-Ranges"] = "bytes" + + range_header = self.headers.get("Range") + if_range = self.headers.get("If-Range") + if not (range_header and srv.ranges) or (if_range and if_range != srv.etag): + self._send(200, payload, head, headers) + return + + start_s, _, end_s = range_header.split("=", 1)[1].partition("-") + start, end = int(start_s), min(int(end_s), len(payload) - 1) + + status = srv.fail_once.pop(start, None) + if status is not None: + self._send(status, b"", head) + return + + with srv._lock: + srv._active += 1 + srv.max_active = max(srv.max_active, srv._active) + try: + if srv.gate_from is not None and start >= srv.gate_from: + srv.gate.wait(timeout=10) + if srv.chunk_delay: + time.sleep(srv.chunk_delay) + body = payload[start:end + 1] + headers["Content-Range"] = f"bytes {start}-{end}/{len(payload)}" + self._send(206, body, head, headers) + if not head and end > 0: + with srv._lock: + srv.ranged_bytes_served += len(body) + finally: + with srv._lock: + srv._active -= 1 + + +def _config(**overrides: object) -> DownloadConfig: + values: dict[str, object] = { + "chunk_size_bytes": MB, + "max_parallel_workers": 4, + "max_retries_per_chunk": 3, + "retry_base_delay_seconds": 0.0, + "retry_jitter_factor": 0.0, + "progress_update_interval_seconds": 0.001, + "http2": False, + } + values.update(overrides) + return DownloadConfig(**values) + + +class DownloadEngineTestBase(unittest.IsolatedAsyncioTestCase): + payload_size = 5 * MB + 12345 + + def setUp(self) -> None: + self.payload = os.urandom(self.payload_size) + self.digest = hashlib.sha256(self.payload).hexdigest() + self.server = _RangeServer(self.payload) + self.thread = threading.Thread(target=self.server.serve_forever, kwargs={"poll_interval": 0.05}, daemon=True) + self.thread.start() + self._tmp = tempfile.TemporaryDirectory() + self.tmp = Path(self._tmp.name) + self.output = self.tmp / "out.bin" + + def tearDown(self) -> None: + self.server.gate.set() + self.server.shutdown() + self.server.server_close() + self._tmp.cleanup() + + @property + def state_path(self) -> Path: + return StateManager().get_state_path(self.output) + + +class TestDownload(DownloadEngineTestBase): + async def test_parallel_download_matches_source(self) -> None: + self.server.chunk_delay = 0.05 + result = await DownloadEngine(_config()).download(self.server.url, self.output) + + self.assertEqual(self.output.read_bytes(), self.payload) + self.assertEqual(result.file_hash, self.digest) + self.assertEqual(result.total_chunks, 6) + self.assertEqual(result.total_bytes_downloaded, self.payload_size) + self.assertGreaterEqual(self.server.max_active, 2, "chunks were not fetched concurrently") + self.assertLessEqual(self.server.max_active, 4, "worker pool exceeded its bound") + self.assertFalse(self.state_path.exists()) + self.assertFalse(self.state_path.parent.exists()) + self.assertFalse(Path(str(self.output) + PART_SUFFIX).exists()) + + async def test_expected_hash_is_verified(self) -> None: + result = await DownloadEngine(_config()).download( + self.server.url, self.output, expected_hash=f"sha256:{self.digest.upper()}" + ) + self.assertTrue(result.is_verified) + + async def test_expected_hash_mismatch_keeps_partial_state(self) -> None: + with self.assertRaises(FileHashMismatchError): + await DownloadEngine(_config()).download(self.server.url, self.output, expected_hash="00" * 32) + self.assertFalse(self.output.exists()) + self.assertEqual(StateManager().load(self.state_path).status, DownloadStatus.FAILED) + + async def test_probe_falls_back_to_range_get_when_head_is_rejected(self) -> None: + self.server.allow_head = False + await DownloadEngine(_config()).download(self.server.url, self.output) + self.assertEqual(self.output.read_bytes(), self.payload) + + async def test_server_without_ranges_is_rejected(self) -> None: + self.server.ranges = False + with self.assertRaises(RangeNotSupportedError): + await DownloadEngine(_config()).download(self.server.url, self.output) + self.assertFalse(self.output.exists()) + + async def test_transient_server_error_is_retried(self) -> None: + self.server.fail_once = {2 * MB: 503, 4 * MB: 429} + result = await DownloadEngine(_config()).download(self.server.url, self.output) + self.assertEqual(self.output.read_bytes(), self.payload) + self.assertEqual(result.chunks_retried, 2) + + async def test_permanent_client_error_aborts_and_checkpoints(self) -> None: + self.server.fail_once = {3 * MB: 404} + with self.assertRaises(ClientError): + await DownloadEngine(_config()).download(self.server.url, self.output) + + state = StateManager().load(self.state_path) + self.assertEqual(state.status, DownloadStatus.FAILED) + self.assertEqual(state.chunks[3].status, ChunkStatus.ABANDONED) + self.assertNotIn(ChunkStatus.IN_PROGRESS, {c.status for c in state.chunks}) + + # The failure was transient after all; resume finishes the job. + result = await DownloadEngine(_config()).resume(self.state_path) + self.assertEqual(self.output.read_bytes(), self.payload) + self.assertEqual(result.file_hash, self.digest) + + async def test_progress_reports_reach_completion(self) -> None: + reports: list[ProgressReport] = [] + await DownloadEngine(_config(), progress_callback=reports.append).download(self.server.url, self.output) + self.assertTrue(reports) + self.assertEqual(reports[-1].downloaded_bytes, self.payload_size) + self.assertEqual(reports[-1].chunks_complete, 6) + percentages = [r.percentage for r in reports] + self.assertEqual(percentages, sorted(percentages)) + + +class TestShutdownAndResume(DownloadEngineTestBase): + async def _interrupt_after_first_chunk(self) -> None: + """Start a download, let chunk 0 finish, then request a graceful stop.""" + self.server.gate_from = MB # every chunk but the first stalls + + engine: DownloadEngine + + def on_progress(report: ProgressReport) -> None: + if report.chunks_complete >= 1: + engine.request_shutdown() + + engine = DownloadEngine(_config(), progress_callback=on_progress) + with self.assertRaises(DownloadCancelledError) as ctx: + await asyncio.wait_for(engine.download(self.server.url, self.output), timeout=10) + self.assertEqual(ctx.exception.state_file, str(self.state_path)) + + async def test_shutdown_checkpoints_completed_chunks(self) -> None: + await self._interrupt_after_first_chunk() + + state = StateManager().load(self.state_path) + self.assertEqual(state.status, DownloadStatus.CANCELLED) + self.assertEqual(state.chunks[0].status, ChunkStatus.COMPLETE) + self.assertEqual(state.chunks[0].computed_hash, hashlib.sha256(self.payload[:MB]).hexdigest()) + self.assertTrue(all(c.status == ChunkStatus.PENDING for c in state.chunks[1:])) + self.assertFalse(self.output.exists()) + + async def test_resume_fetches_only_missing_chunks(self) -> None: + await self._interrupt_after_first_chunk() + self.server.gate.set() + self.server.gate_from = None + served_before = self.server.ranged_bytes_served + + result = await DownloadEngine(_config()).resume(self.state_path) + + self.assertEqual(self.output.read_bytes(), self.payload) + self.assertEqual(result.total_bytes_downloaded, self.payload_size - MB) + self.assertEqual(self.server.ranged_bytes_served - served_before, self.payload_size - MB) + self.assertFalse(self.state_path.exists()) + + async def test_resume_refetches_chunk_corrupted_on_disk(self) -> None: + await self._interrupt_after_first_chunk() + self.server.gate.set() + self.server.gate_from = None + part = Path(str(self.output) + PART_SUFFIX) + with open(part, "r+b") as f: + f.seek(100) + f.write(b"\x00" * 16) + + result = await DownloadEngine(_config()).resume(self.state_path) + + self.assertEqual(self.output.read_bytes(), self.payload) + self.assertEqual(result.total_bytes_downloaded, self.payload_size) + + async def test_resume_rejects_changed_remote_file(self) -> None: + await self._interrupt_after_first_chunk() + self.server.gate.set() + self.server.gate_from = None + self.server.etag = '"v2"' + + with self.assertRaises(PreconditionFailedError): + await DownloadEngine(_config()).resume(self.state_path) + self.assertFalse(self.output.exists()) + + async def test_cancelling_the_task_checkpoints_state(self) -> None: + self.server.gate_from = MB + task = asyncio.create_task(DownloadEngine(_config()).download(self.server.url, self.output)) + while not self.state_path.exists() or StateManager().load(self.state_path).chunks[0].status != ChunkStatus.COMPLETE: + await asyncio.sleep(0.02) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + state = StateManager().load(self.state_path) + self.assertEqual(state.status, DownloadStatus.CANCELLED) + self.assertEqual(state.chunks[0].status, ChunkStatus.COMPLETE) + + +class TestCli(DownloadEngineTestBase): + payload_size = 2 * MB + 7 + + def _run(self, *argv: str) -> tuple[int, str]: + out = io.StringIO() + with redirect_stdout(out): + code = cli_main(list(argv)) + return code, out.getvalue() + + def test_download_command(self) -> None: + code, out = self._run( + "download", "--url", self.server.url, "-o", str(self.output), + "--workers", "2", "--chunk-size", "1MB", "--sha256", self.digest, + ) + self.assertEqual(code, 0, out) + self.assertIn("[SUCCESS]", out) + self.assertIn(f"SHA-256 (verified): {self.digest}", out) + self.assertEqual(self.output.read_bytes(), self.payload) + + def test_download_command_reports_errors(self) -> None: + self.server.ranges = False + code, out = self._run("download", "--url", self.server.url, "-o", str(self.output)) + self.assertEqual(code, 1) + self.assertIn("[ERROR]", out) + + def test_invalid_chunk_size_is_rejected(self) -> None: + code, out = self._run("download", "--url", self.server.url, "-o", str(self.output), "--chunk-size", "1KB") + self.assertEqual(code, 1) + self.assertIn("Invalid configuration", out) + + def test_resume_command(self) -> None: + self.server.fail_once = {MB: 404} + code, _ = self._run("download", "--url", self.server.url, "-o", str(self.output), "--chunk-size", "1MB") + self.assertEqual(code, 1) + code, out = self._run("resume", "--state-file", str(self.state_path)) + self.assertEqual(code, 0, out) + self.assertEqual(self.output.read_bytes(), self.payload) + + +class TestHelpers(unittest.TestCase): + def test_plan_chunks_covers_file_without_gaps(self) -> None: + chunks = plan_chunks(10 * MB + 1, 4 * MB) + self.assertEqual([(c.start_byte, c.end_byte) for c in chunks], [ + (0, 4 * MB - 1), (4 * MB, 8 * MB - 1), (8 * MB, 10 * MB), + ]) + self.assertEqual(sum(c.size for c in chunks), 10 * MB + 1) + + def test_plan_chunks_exact_multiple_and_empty(self) -> None: + self.assertEqual(len(plan_chunks(8 * MB, 4 * MB)), 2) + self.assertEqual(plan_chunks(0, 4 * MB), []) + + def test_parse_content_range(self) -> None: + self.assertEqual(_parse_content_range("bytes 0-0/1234"), (0, 0, 1234)) + self.assertEqual(_parse_content_range("bytes 10-19/*"), (10, 19, None)) + self.assertIsNone(_parse_content_range("items 0-1/2")) + self.assertIsNone(_parse_content_range(None)) + + +if __name__ == "__main__": + unittest.main()