tests/test_benchmark_source.py

from __future__ import annotations

import json
from pathlib import Path
from types import SimpleNamespace
from typing import cast

import pytest

import harbor_hf.benchmark_source as benchmark_source
from harbor_hf.benchmark_source import (
    BenchmarkSourceLock,
    PackageBenchmarkSourceLock,
    anonymous_git_environment,
    bundle_prefix,
    bundle_uri,
    known_credential_values,
    load_source_lock,
    prepare_benchmark_source,
    resolve_benchmark_source,
    resolved_experiment,
    source_lock_bytes,
    source_lock_digest,
    source_lock_from_spec,
    source_lock_json_schema,
    verify_anonymous_git_source,
)
from harbor_hf.control import RunSubmittedPayload, new_event
from harbor_hf.executions import build_execution_lock
from harbor_hf.harbor_adapter import build_execution_request
from harbor_hf.models import (
    BundleBenchmarkSource,
    DirectoryBenchmarkSource,
    ExperimentSpec,
    GitBenchmarkSource,
)
from harbor_hf.reconciler import plan_reconciliation
from harbor_hf.runs import (
    RunLock,
    build_run_lock,
    build_run_plan,
    build_wave_lock,
)


def _with_source(spec: ExperimentSpec, source: object) -> ExperimentSpec:
    raw = spec.model_dump(mode="python")
    raw["benchmark"]["dataset"] = "shellbench/local"
    raw["benchmark"]["source"] = source
    raw["benchmark"].pop("dataset_digest", None)
    return ExperimentSpec.model_validate(raw)


def test_directory_resolution_is_content_addressed_and_path_independent(
    remote_spec: ExperimentSpec,
    tmp_path: Path,
) -> None:
    first = tmp_path / "first" / "tasks"
    second = tmp_path / "second" / "tasks"
    first.mkdir(parents=True)
    second.mkdir(parents=True)
    (first / "task.toml").write_text("name='same'\n", encoding="utf-8")
    (second / "task.toml").write_text("name='same'\n", encoding="utf-8")
    first_manifest = tmp_path / "first" / "run.yaml"
    second_manifest = tmp_path / "second" / "run.yaml"
    first_manifest.write_text("run\n", encoding="utf-8")
    second_manifest.write_text("run\n", encoding="utf-8")
    spec = _with_source(
        remote_spec,
        DirectoryBenchmarkSource(path="tasks").model_dump(mode="python"),
    )

    first_resolution = resolve_benchmark_source(
        spec, first_manifest, tmp_path / "workspace-one"
    )
    second_resolution = resolve_benchmark_source(
        spec, second_manifest, tmp_path / "workspace-two"
    )

    assert first_resolution.lock == second_resolution.lock
    assert source_lock_digest(first_resolution.lock) == source_lock_digest(
        second_resolution.lock
    )
    assert first_resolution.bundle is not None
    assert second_resolution.bundle is not None
    assert "first" not in source_lock_bytes(first_resolution.lock).decode()
    assert "second" not in source_lock_bytes(second_resolution.lock).decode()
    resolved = resolved_experiment(spec, first_resolution.lock)
    assert isinstance(resolved.benchmark.source, BundleBenchmarkSource)
    assert resolved.benchmark.dataset_digest == resolved.benchmark.source.content_digest


@pytest.mark.parametrize(
    "path",
    ["/absolute", ".", "tasks/../private", "tasks\\private", "tasks\0private"],
)
def test_git_source_rejects_unsafe_paths(path: str) -> None:
    with pytest.raises(ValueError, match="safely relative"):
        GitBenchmarkSource(
            repository="ShellBench/public-tasks",
            revision="8" * 40,
            path=path,
        )


def test_source_lock_is_canonical_and_supports_package_git_and_bundle(
    remote_spec: ExperimentSpec,
    tmp_path: Path,
) -> None:
    package = source_lock_from_spec(remote_spec)
    assert isinstance(package.source, PackageBenchmarkSourceLock)
    assert package.source.reference == "harbor/terminal-bench@sha256:" + "1" * 64

    git = GitBenchmarkSource(
        repository="https://github.com/ShellBench/public-tasks.git",
        revision="8" * 40,
        path="tasks/115-tasks",
    )
    git_spec = _with_source(remote_spec, git.model_dump(mode="python"))
    git_lock = resolve_benchmark_source(
        git_spec,
        tmp_path / "run.yaml",
        tmp_path / "git-workspace",
        verify_git=False,
    ).lock
    assert git_lock.source == git

    bundle = BundleBenchmarkSource(
        content_digest="sha256:" + "2" * 64,
        manifest_sha256="3" * 64,
    )
    bundle_spec = _with_source(remote_spec, bundle.model_dump(mode="python"))
    assert source_lock_from_spec(bundle_spec).source == bundle
    assert (
        bundle_prefix(bundle.content_digest) == "benchmark-bundles/sha256/" + "2" * 64
    )
    assert bundle_uri("example-org", bundle.content_digest) == (
        "hf://buckets/example-org/jobs-artifacts/benchmark-bundles/sha256/" + "2" * 64
    )

    path = tmp_path / "source.lock.json"
    path.write_bytes(source_lock_bytes(git_lock))
    assert load_source_lock(path) == git_lock
    value = json.loads(path.read_text(encoding="utf-8"))
    path.write_text(json.dumps(value, indent=2), encoding="utf-8")
    with pytest.raises(ValueError, match="canonical JSON"):
        load_source_lock(path)


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


def test_source_identity_changes_run_shard_trial_and_wave_ids(
    remote_spec: ExperimentSpec,
) -> None:
    first_source = GitBenchmarkSource(
        repository="ShellBench/public-tasks",
        revision="8" * 40,
        path="tasks/115-tasks",
    )
    second_source = first_source.model_copy(update={"revision": "9" * 40})
    first_spec = _with_source(remote_spec, first_source.model_dump(mode="python"))
    second_spec = _with_source(remote_spec, second_source.model_dump(mode="python"))
    first_lock = build_run_lock(build_run_plan(first_spec), "same-run")
    second_lock = build_run_lock(build_run_plan(second_spec), "same-run")

    first_run = first_lock.executions[0]
    second_run = second_lock.executions[0]
    assert first_run.execution_id != second_run.execution_id
    assert first_run.shards[0].shard_id != second_run.shards[0].shard_id
    assert (
        first_run.shards[0].trials[0].trial_id
        != second_run.shards[0].trials[0].trial_id
    )

    def wave_id(run: RunLock, spec: ExperimentSpec) -> str:
        submitted = new_event(
            subject_type="run",
            subject_id=run.run_id,
            kind="run.submitted",
            producer="cli",
            payload=RunSubmittedPayload(plan_digest=run.plan_digest),
        )
        action = plan_reconciliation(run, [submitted])[1].actions[0]
        return build_wave_lock(run, spec, action).wave_id

    assert wave_id(first_lock, first_spec) != wave_id(second_lock, second_spec)


def test_source_resolution_rejects_a_changed_request(
    remote_spec: ExperimentSpec,
) -> None:
    requested = GitBenchmarkSource(
        repository="ShellBench/public-tasks",
        revision="8" * 40,
        path="tasks/115-tasks",
    )
    spec = _with_source(remote_spec, requested.model_dump(mode="python"))
    changed = BenchmarkSourceLock(
        source=requested.model_copy(update={"revision": "9" * 40})
    )

    with pytest.raises(ValueError, match="changed the request"):
        resolved_experiment(spec, changed)


def test_anonymous_git_environment_removes_ambient_authentication() -> None:
    environment = anonymous_git_environment(
        {
            "PATH": "/bin",
            "GITHUB_TOKEN": "github-secret",
            "GH_TOKEN": "gh-secret",
            "SSH_AUTH_SOCK": "/tmp/agent",
            "GIT_CONFIG_COUNT": "1",
            "GIT_CONFIG_KEY_0": "http.extraHeader",
            "GIT_CONFIG_VALUE_0": "Authorization: secret",
            "GIT_CONFIG_PARAMETERS": "'credential.helper'='bad'",
        }
    )

    assert environment["PATH"] == "/bin"
    assert environment["GITHUB_TOKEN"] == ""
    assert environment["GH_TOKEN"] == ""
    assert environment["SSH_AUTH_SOCK"] == ""
    assert environment["GIT_TERMINAL_PROMPT"] == "0"
    assert environment["GIT_CONFIG_NOSYSTEM"] == "1"
    assert environment["GIT_CONFIG_GLOBAL"] == "/dev/null"
    assert environment["GIT_CONFIG_SYSTEM"] == "/dev/null"
    assert environment["GIT_CONFIG_COUNT"] == "0"
    assert "GIT_CONFIG_KEY_0" not in environment
    assert "GIT_CONFIG_VALUE_0" not in environment
    assert environment["GIT_CONFIG_PARAMETERS"] == ""


def test_anonymous_git_preflight_checks_exact_commit_and_tree(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    source = GitBenchmarkSource(
        repository="ShellBench/public-tasks",
        revision="8" * 40,
        path="tasks/115-tasks",
    )
    calls: list[tuple[list[str], dict[str, str]]] = []
    results = iter(
        [
            SimpleNamespace(returncode=0, stdout="", stderr=""),
            SimpleNamespace(returncode=0, stdout="", stderr=""),
            SimpleNamespace(returncode=0, stdout=source.revision + "\n", stderr=""),
            SimpleNamespace(returncode=0, stdout="tree\n", stderr=""),
            SimpleNamespace(
                returncode=0,
                stdout="100644 blob 0123456789abcdef\ttasks/115-tasks/task.toml\n",
                stderr="",
            ),
        ]
    )

    def run(command: list[str], **kwargs: object) -> object:
        environment = cast(dict[str, str], kwargs["env"])
        calls.append((command, environment))
        return next(results)

    monkeypatch.setenv("GITHUB_TOKEN", "must-not-leak")
    monkeypatch.setenv("GIT_CONFIG_COUNT", "1")
    monkeypatch.setenv("GIT_CONFIG_KEY_0", "http.extraHeader")
    monkeypatch.setenv("GIT_CONFIG_VALUE_0", "Authorization: must-not-leak")
    monkeypatch.setattr(benchmark_source.subprocess, "run", run)

    verify_anonymous_git_source(source)

    assert len(calls) == 5
    assert calls[1][0][-2:] == [
        "https://github.com/ShellBench/public-tasks.git",
        source.revision,
    ]
    assert calls[2][0][-2:] == ["rev-parse", "FETCH_HEAD"]
    assert calls[3][0][-2:] == ["-t", "FETCH_HEAD:tasks/115-tasks"]
    assert calls[4][0][-2:] == ["--", "tasks/115-tasks"]
    for _command, environment in calls:
        assert "must-not-leak" not in repr(environment)
        assert environment["GIT_CONFIG_COUNT"] == "0"


def test_anonymous_git_preflight_rejects_a_missing_tree(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    source = GitBenchmarkSource(
        repository="ShellBench/public-tasks",
        revision="8" * 40,
        path="tasks/missing",
    )
    results = iter(
        [
            SimpleNamespace(returncode=0, stdout="", stderr=""),
            SimpleNamespace(returncode=0, stdout="", stderr=""),
            SimpleNamespace(returncode=0, stdout=source.revision + "\n", stderr=""),
            SimpleNamespace(returncode=1, stdout="", stderr="missing"),
            SimpleNamespace(returncode=0, stdout="", stderr=""),
        ]
    )
    monkeypatch.setattr(
        benchmark_source.subprocess,
        "run",
        lambda *_args, **_kwargs: next(results),
    )

    with pytest.raises(ValueError, match="path is not a directory"):
        verify_anonymous_git_source(source)


def test_anonymous_git_preflight_rejects_submodules(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    source = GitBenchmarkSource(
        repository="ShellBench/public-tasks",
        revision="8" * 40,
        path="tasks/115-tasks",
    )
    results = iter(
        [
            SimpleNamespace(returncode=0, stdout="", stderr=""),
            SimpleNamespace(returncode=0, stdout="", stderr=""),
            SimpleNamespace(returncode=0, stdout=source.revision + "\n", stderr=""),
            SimpleNamespace(returncode=0, stdout="tree\n", stderr=""),
            SimpleNamespace(
                returncode=0,
                stdout=(
                    "160000 commit 0123456789abcdef"
                    "\ttasks/115-tasks/private-dependency\n"
                ),
                stderr="",
            ),
        ]
    )
    monkeypatch.setattr(
        benchmark_source.subprocess,
        "run",
        lambda *_args, **_kwargs: next(results),
    )

    with pytest.raises(ValueError, match="cannot contain submodules"):
        verify_anonymous_git_source(source)


def test_bundle_loading_extracts_and_renders_a_local_harbor_dataset(
    remote_spec: ExperimentSpec,
    tmp_path: Path,
) -> None:
    source = tmp_path / "tasks"
    source.mkdir()
    (source / "task.toml").write_text("name='local'\n", encoding="utf-8")
    spec = _with_source(
        remote_spec,
        DirectoryBenchmarkSource(path=str(source)).model_dump(mode="python"),
    )
    resolution = resolve_benchmark_source(
        spec, tmp_path / "run.yaml", tmp_path / "workspace"
    )
    assert resolution.bundle is not None
    extracted = prepare_benchmark_source(
        resolution.lock,
        mounted_bundle_root=resolution.bundle.bundle_root,
        destination=tmp_path / "extracted",
    )
    assert extracted == tmp_path / "extracted"
    assert extracted is not None
    resolved = resolved_experiment(spec, resolution.lock)
    lock = build_execution_lock(resolved, execution_id="bundle-run")
    request = build_execution_request(
        lock,
        tmp_path / "jobs",
        "https://endpoint.example",
        task_names=lock.benchmark_tasks,
        attempts=1,
        concurrency=1,
        expected_task_digests=lock.benchmark_task_digests,
        benchmark_root=extracted,
    )

    datasets = request.harbor_config["datasets"]
    assert isinstance(datasets, list)
    assert datasets == [
        {"path": str(extracted.resolve()), "task_names": lock.benchmark_tasks}
    ]


def test_remote_git_preparation_repeats_anonymous_preflight(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    source = GitBenchmarkSource(
        repository="ShellBench/public-tasks",
        revision="8" * 40,
        path="tasks/115-tasks",
    )
    observed: list[GitBenchmarkSource] = []
    monkeypatch.setattr(
        benchmark_source,
        "verify_anonymous_git_source",
        observed.append,
    )

    result = prepare_benchmark_source(
        BenchmarkSourceLock(source=source),
        mounted_bundle_root=tmp_path / "unused",
        destination=tmp_path / "unused-destination",
    )

    assert result is None
    assert observed == [source]


def test_known_credential_values_selects_only_credential_like_environment() -> None:
    assert known_credential_values(
        {
            "PATH": "/ordinary/path",
            "SHORT_TOKEN": "short",
            "HARBOR_HF_JOB_TOKEN": "purpose-scoped-value",
            "SERVICE_PASSWORD": "password-value",
            "UNRELATED": "not-selected",
        }
    ) == ("purpose-scoped-value", "password-value")