test_ranges.py

#!/usr/bin/env python3
"""Compare random byte ranges from DNA Zoo Wasabi and Hugging Face Buckets.

Uses only the Python standard library. The full 2.52 GB file is never
downloaded; each endpoint is asked for a small, independently selected range.
"""

from __future__ import annotations

import argparse
import hashlib
import random
import re
import subprocess
import tempfile
import time


SOURCE_URL = (
    "https://dnazoo.s3.wasabisys.com/"
    "Acinonyx_jubatus/aciJub1.hic"
)
HF_URL = (
    "https://huggingface.co/buckets/abidlabs/dnazoo-poc/resolve/"
    "Acinonyx_jubatus/aciJub1.hic"
)
FILE_SIZE = 2_524_603_208
TEST_ORIGIN = "https://www.dnazoo.org"


def fetch_range(url: str, start: int, end: int, *, origin: str | None = None):
    with tempfile.TemporaryDirectory(prefix="dnazoo-range-") as temp_dir:
        body_path = f"{temp_dir}/body.bin"
        headers_path = f"{temp_dir}/headers.txt"
        command = [
            "curl",
            "--silent",
            "--show-error",
            "--location",
            "--fail",
            "--max-time",
            "60",
            "--user-agent",
            "dnazoo-hf-range-poc/1.0",
            "--header",
            f"Range: bytes={start}-{end}",
            "--dump-header",
            headers_path,
            "--output",
            body_path,
            "--write-out",
            "%{http_code}\n%{url_effective}\n%{time_total}",
        ]
        if origin:
            command.extend(["--header", f"Origin: {origin}"])
        command.append(url)

        started = time.perf_counter()
        result = subprocess.run(command, check=True, capture_output=True, text=True)
        elapsed = time.perf_counter() - started
        status_text, final_url, curl_time = result.stdout.splitlines()
        body = open(body_path, "rb").read()
        raw_headers = open(headers_path, encoding="iso-8859-1").read()

    # `curl -L` records each redirect response. Use the final HTTP header block.
    blocks = [
        block for block in re.split(r"\r?\n\r?\n", raw_headers.strip())
        if block.startswith("HTTP/")
    ]
    final_headers = {}
    for line in blocks[-1].splitlines()[1:]:
        if ":" in line:
            key, value = line.split(":", 1)
            final_headers[key.strip().lower()] = value.strip()

    metadata = {
        "status": int(status_text),
        "content_range": final_headers.get("content-range"),
        "accept_ranges": final_headers.get("accept-ranges"),
        "cors": final_headers.get("access-control-allow-origin"),
        "final_url": final_url,
        "curl_time": float(curl_time),
    }
    return body, metadata, elapsed


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Verify random-access equivalence between Wasabi and HF Buckets."
    )
    parser.add_argument("--samples", type=int, default=7)
    parser.add_argument("--bytes", type=int, default=65_536, dest="range_size")
    args = parser.parse_args()

    if args.samples < 3:
        parser.error("--samples must be at least 3")
    if not 1 <= args.range_size <= FILE_SIZE:
        parser.error(f"--bytes must be between 1 and {FILE_SIZE}")

    # Always cover the header, a deterministic set of scattered positions,
    # and the final bytes where a format may keep indexes or metadata.
    starts = [0]
    rng = random.Random(20260821)
    for _ in range(args.samples - 2):
        starts.append(rng.randrange(1, FILE_SIZE - args.range_size))
    starts.append(FILE_SIZE - args.range_size)

    print(f"File size: {FILE_SIZE:,} bytes")
    print(f"Range size: {args.range_size:,} bytes")
    print(f"Samples:    {len(starts)}\n")
    print("#  range                           Wasabi     HF       SHA-256")
    print("-  ------------------------------  ---------  -------  ----------------")

    for index, start in enumerate(starts, 1):
        end = start + args.range_size - 1
        source, source_meta, source_time = fetch_range(SOURCE_URL, start, end)
        target, target_meta, target_time = fetch_range(
            HF_URL, start, end, origin=TEST_ORIGIN
        )

        expected_content_range = f"bytes {start}-{end}/{FILE_SIZE}"
        for name, body, metadata in (
            ("Wasabi", source, source_meta),
            ("HF", target, target_meta),
        ):
            if metadata["status"] != 206:
                raise AssertionError(f"{name}: expected HTTP 206, got {metadata['status']}")
            if metadata["content_range"] != expected_content_range:
                raise AssertionError(
                    f"{name}: expected {expected_content_range!r}, "
                    f"got {metadata['content_range']!r}"
                )
            if len(body) != args.range_size:
                raise AssertionError(
                    f"{name}: expected {args.range_size} bytes, got {len(body)}"
                )

        if target_meta["accept_ranges"] != "bytes":
            raise AssertionError(
                f"HF: expected Accept-Ranges: bytes, got {target_meta['accept_ranges']!r}"
            )
        if target_meta["cors"] not in ("*", TEST_ORIGIN):
            raise AssertionError(
                f"HF: CORS does not allow {TEST_ORIGIN}: {target_meta['cors']!r}"
            )
        if source != target:
            raise AssertionError(f"byte mismatch for range {start}-{end}")

        digest = hashlib.sha256(target).hexdigest()[:16]
        print(
            f"{index:<2} {start:>13,}-{end:<13,}  "
            f"{source_time:>7.3f}s  {target_time:>5.3f}s  {digest}"
        )

    print("\nPASS: all ranges returned HTTP 206 and matched byte-for-byte.")
    print(f"PASS: HF response allows browser origin {TEST_ORIGIN}.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())