tests/test_benchmark_bundle.py

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path

import pytest
import zstandard

import harbor_hf.benchmark_bundle as benchmark_bundle
from harbor_hf.benchmark_bundle import (
    BenchmarkBundleManifest,
    BundleDirectory,
    BundleEntry,
    BundleFile,
    BundlePayload,
    benchmark_bundle_json_schema,
    build_benchmark_bundle,
    bundle_content_digest,
    bundle_manifest_bytes,
    extract_benchmark_bundle,
    validate_benchmark_bundle,
    verify_benchmark_bundle,
)


def _source(root: Path) -> Path:
    source = root / "tasks"
    (source / "nested").mkdir(parents=True)
    (source / "dataset.toml").write_text("name = 'local'\n", encoding="utf-8")
    executable = source / "nested" / "verify.sh"
    executable.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8")
    executable.chmod(0o751)
    return source


def test_bundle_is_deterministic_and_extracts_exact_contents(tmp_path: Path) -> None:
    first_source = _source(tmp_path / "first")
    second_source = _source(tmp_path / "second")
    first = build_benchmark_bundle(first_source, tmp_path / "bundle-one")
    second = build_benchmark_bundle(second_source, tmp_path / "bundle-two")

    assert first.manifest == second.manifest
    assert first.manifest_sha256 == second.manifest_sha256
    assert first.payload_path.read_bytes() == second.payload_path.read_bytes()
    assert str(first_source) not in first.manifest_path.read_text(encoding="utf-8")

    destination = tmp_path / "extracted"
    previous_umask = os.umask(0o077)
    try:
        assert extract_benchmark_bundle(first.bundle_root, destination) == destination
    finally:
        os.umask(previous_umask)
    assert (destination / "dataset.toml").read_bytes() == b"name = 'local'\n"
    assert (destination / "nested" / "verify.sh").read_bytes() == (
        b"#!/bin/sh\nexit 0\n"
    )
    assert (destination / "nested").stat().st_mode & 0o777 == 0o755
    assert (destination / "dataset.toml").stat().st_mode & 0o777 == 0o644
    assert (destination / "nested" / "verify.sh").stat().st_mode & 0o777 == 0o755
    assert verify_benchmark_bundle(first.bundle_root) == first.manifest


def test_bundle_identity_changes_with_bytes_modes_and_paths(tmp_path: Path) -> None:
    source = tmp_path / "source"
    source.mkdir()
    file = source / "task.toml"
    file.write_text("a", encoding="utf-8")
    initial = build_benchmark_bundle(source, tmp_path / "initial")

    file.write_text("b", encoding="utf-8")
    changed_bytes = build_benchmark_bundle(source, tmp_path / "changed-bytes")
    file.chmod(0o755)
    changed_mode = build_benchmark_bundle(source, tmp_path / "changed-mode")
    file.rename(source / "renamed.toml")
    changed_path = build_benchmark_bundle(source, tmp_path / "changed-path")

    digests = {
        bundle.manifest.content_digest
        for bundle in (initial, changed_bytes, changed_mode, changed_path)
    }
    assert len(digests) == 4


def test_bundle_rejects_unsafe_source_nodes_and_destinations(tmp_path: Path) -> None:
    source = _source(tmp_path)
    (source / "link").symlink_to(source / "dataset.toml")
    with pytest.raises(ValueError, match="symbolic links"):
        build_benchmark_bundle(source, tmp_path / "link-bundle")
    (source / "link").unlink()

    (source / ".git").mkdir()
    with pytest.raises(ValueError, match="safe normalized relative path"):
        build_benchmark_bundle(source, tmp_path / "git-bundle")
    (source / ".git").rmdir()

    if hasattr(os, "mkfifo"):
        os.mkfifo(source / "pipe")
        with pytest.raises(ValueError, match="special file"):
            build_benchmark_bundle(source, tmp_path / "pipe-bundle")
        (source / "pipe").unlink()

    invalid_path = os.fsencode(source) + b"/invalid-\xff"
    descriptor = os.open(invalid_path, os.O_CREAT | os.O_WRONLY, 0o600)
    os.close(descriptor)
    with pytest.raises(ValueError, match="not valid UTF-8"):
        build_benchmark_bundle(source, tmp_path / "invalid-name-bundle")
    os.unlink(invalid_path)

    with pytest.raises(ValueError, match="outside its source"):
        build_benchmark_bundle(source, source / "generated")


def test_bundle_rejects_source_mutation_during_construction(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    source = _source(tmp_path)
    original = benchmark_bundle._scan_source
    calls = 0

    def scan(
        root: Path, *, snapshot: Path | None, needles: tuple[bytes, ...]
    ) -> list[BundleEntry]:
        nonlocal calls
        entries = original(root, snapshot=snapshot, needles=needles)
        calls += 1
        if calls == 1:
            (source / "dataset.toml").write_text("changed\n", encoding="utf-8")
        return entries

    monkeypatch.setattr(benchmark_bundle, "_scan_source", scan)

    with pytest.raises(ValueError, match="changed while its bundle was built"):
        build_benchmark_bundle(source, tmp_path / "mutating-bundle")


def test_bundle_rejects_known_credentials_and_private_keys(tmp_path: Path) -> None:
    source = _source(tmp_path)
    secret = "purpose-scoped-secret-123"
    boundary = source / "boundary.bin"
    boundary.write_bytes(b"x" * (1024 * 1024 - 5) + secret.encode())
    with pytest.raises(ValueError, match="credential material"):
        build_benchmark_bundle(
            source,
            tmp_path / "secret-bundle",
            known_secrets=(secret,),
        )

    boundary.write_text(
        "-----BEGIN OPENSSH PRIVATE KEY-----\nnot-a-real-key\n",
        encoding="utf-8",
    )
    with pytest.raises(ValueError, match="credential material"):
        build_benchmark_bundle(source, tmp_path / "key-bundle")


def test_bundle_validation_rejects_changed_or_noncanonical_files(
    tmp_path: Path,
) -> None:
    prepared = build_benchmark_bundle(_source(tmp_path), tmp_path / "bundle")
    prepared.payload_path.write_bytes(prepared.payload_path.read_bytes() + b"changed")
    with pytest.raises(ValueError, match="byte count"):
        validate_benchmark_bundle(prepared.bundle_root)

    prepared = build_benchmark_bundle(
        _source(tmp_path / "again"), tmp_path / "bundle-two"
    )
    manifest = json.loads(prepared.manifest_path.read_text(encoding="utf-8"))
    prepared.manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    with pytest.raises(ValueError, match="canonical JSON"):
        validate_benchmark_bundle(prepared.bundle_root)

    (prepared.bundle_root / "extra").write_text("x", encoding="utf-8")
    with pytest.raises(ValueError, match="exactly its manifest and payload"):
        validate_benchmark_bundle(prepared.bundle_root)


def test_bundle_manifest_requires_sorted_complete_parent_inventory() -> None:
    file_entry = BundleFile(
        path="nested/file.txt",
        mode=420,
        bytes=1,
        sha256="0" * 64,
    )
    with pytest.raises(ValueError, match="missing a parent directory"):
        BenchmarkBundleManifest.model_validate(
            {
                "content_digest": bundle_content_digest([file_entry]),
                "entries": [file_entry.model_dump(mode="json")],
                "payload": {"bytes": 1, "sha256": "1" * 64},
            }
        )

    entries = [
        file_entry,
        BundleDirectory(path="nested"),
    ]
    with pytest.raises(ValueError, match="sorted by UTF-8 path"):
        BenchmarkBundleManifest.model_validate(
            {
                "content_digest": bundle_content_digest(entries),
                "entries": [entry.model_dump(mode="json") for entry in entries],
                "payload": {"bytes": 1, "sha256": "1" * 64},
            }
        )


def test_bundle_extraction_rejects_a_decompression_overrun(tmp_path: Path) -> None:
    root = tmp_path / "bomb"
    root.mkdir()
    payload = zstandard.ZstdCompressor().compress(b"x" * (2 * 1024 * 1024))
    (root / "payload.tar.zst").write_bytes(payload)
    entry = BundleFile(
        path="empty",
        mode=420,
        bytes=0,
        sha256=hashlib.sha256(b"").hexdigest(),
    )
    manifest = BenchmarkBundleManifest(
        content_digest=bundle_content_digest([entry]),
        entries=[entry],
        payload=BundlePayload(
            bytes=len(payload), sha256=hashlib.sha256(payload).hexdigest()
        ),
    )
    (root / "bundle.json").write_bytes(bundle_manifest_bytes(manifest))

    with pytest.raises(ValueError, match="decompressed past its limit"):
        extract_benchmark_bundle(root, tmp_path / "bomb-output")


def test_checked_in_bundle_schema_matches_the_model() -> None:
    path = (
        Path(__file__).parents[1] / "schemas" / "benchmark-bundle-v1alpha1.schema.json"
    )
    assert (
        json.loads(path.read_text(encoding="utf-8")) == benchmark_bundle_json_schema()
    )


def test_bundle_extraction_requires_an_empty_destination(tmp_path: Path) -> None:
    prepared = build_benchmark_bundle(_source(tmp_path), tmp_path / "bundle")
    destination = tmp_path / "destination"
    destination.mkdir()
    (destination / "existing").write_text("x", encoding="utf-8")

    with pytest.raises(ValueError, match="destination must be empty"):
        extract_benchmark_bundle(prepared.bundle_root, destination)