giant/tests/test_dwarf.py

from typer.testing import CliRunner

from giant import cli as giant_cli
from giant.config import Conditioning
from giant.data import setup_cache
from scripts import dwarf
from scripts.dwarf import app
from test_pipeline import _make_synthetic_steps

runner = CliRunner()


def test_conditioning_enum_shared_across_both_clis():
    """giant.cli and scripts.dwarf must use the one giant.config.Conditioning
    enum, not independently redefined copies that could silently drift apart
    on valid --conditioning values."""
    assert dwarf.Conditioning is Conditioning
    assert giant_cli.Conditioning is Conditioning


def test_warn_if_exceeds_shared_quota_warns_over_quarter_cpu(monkeypatch, capsys):
    monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8)  # quota = 2
    dwarf._warn_if_exceeds_shared_quota(3, "--jobs")
    assert "warning: --jobs=3 exceeds" in capsys.readouterr().err


def test_warn_if_exceeds_shared_quota_silent_within_quota(monkeypatch, capsys):
    monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8)  # quota = 2
    dwarf._warn_if_exceeds_shared_quota(2, "--jobs")
    assert capsys.readouterr().err == ""


def test_convert_rejects_jobs_below_one(tmp_path):
    root_file = tmp_path / "shard.root"
    root_file.touch()
    result = runner.invoke(app, ["convert", str(root_file), "--jobs", "0"])
    assert result.exit_code != 0
    assert "--jobs must be >= 1" in result.output


def test_convert_rejects_output_with_multiple_files(tmp_path):
    a = tmp_path / "a.root"
    b = tmp_path / "b.root"
    a.touch()
    b.touch()
    result = runner.invoke(app, ["convert", str(a), str(b), "--output", "out.parquet"])
    assert result.exit_code != 0
    assert "--output can only be used with a single input file" in result.output


def test_convert_rejects_output_with_parallel_jobs(tmp_path):
    root_file = tmp_path / "shard.root"
    root_file.touch()
    result = runner.invoke(
        app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
    )
    assert result.exit_code != 0
    assert "--output cannot be combined with --jobs > 1" in result.output


def test_convert_default_jobs_is_one():
    result = runner.invoke(app, ["convert", "--help"])
    assert result.exit_code == 0
    assert "default: 1" in result.output


def test_bump_gen_requires_reason():
    result = runner.invoke(app, ["bump-gen"])
    assert result.exit_code != 0
    assert "reason" in result.output.lower()


def test_create_manifest_requires_exactly_one_of_output_or_pool(tmp_path):
    f = tmp_path / "a.parquet"
    f.touch()
    result = runner.invoke(app, ["create-manifest", str(f)])
    assert result.exit_code != 0
    assert "exactly one of --output or --pool is required" in result.output


def test_create_manifest_requires_type_with_pool(tmp_path):
    f = tmp_path / "a.parquet"
    f.touch()
    result = runner.invoke(app, ["create-manifest", "--pool", "pbwo4", str(f)])
    assert result.exit_code != 0
    assert "--type is required when --pool is given" in result.output


def test_status_reports_missing_root(tmp_path):
    missing = tmp_path / "does-not-exist"
    result = runner.invoke(app, ["status", "--root", str(missing)])
    assert result.exit_code != 0
    assert "is not a directory" in result.output


def test_warm_cache_writes_sidecar(tmp_path):
    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)

    result = runner.invoke(app, ["warm-cache", str(data)])

    assert result.exit_code == 0, result.output
    loaded = setup_cache.load(data, [data])
    assert loaded is not None
    assert loaded.vocab is not None
    assert loaded.event_index is not None
    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers


def test_warm_cache_second_run_hits_cache(tmp_path):
    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
    runner.invoke(app, ["warm-cache", str(data)])

    result = runner.invoke(app, ["warm-cache", str(data)])

    assert result.exit_code == 0, result.output
    assert "event index: cache hit" in result.output
    assert "vocabulary maps: cache hit" in result.output
    assert "normalizer: cache hit" in result.output


def test_warm_cache_router_process_warms_proc_map(tmp_path):
    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)

    result = runner.invoke(
        app,
        [
            "warm-cache",
            str(data),
            "--router",
            "--router-type",
            "process",
            "--n-experts",
            "3",
        ],
    )

    assert result.exit_code == 0, result.output
    loaded = setup_cache.load(data, [data])
    assert loaded is not None
    assert 3 in loaded.proc_maps


def test_warm_cache_rebuild_ignores_existing(tmp_path):
    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
    files = [data]
    stale = setup_cache.SetupCache.empty(files)
    stale.vocab = ({999999: 0}, {"G4_AIR": 0})  # deliberately wrong
    setup_cache.save(data, files, stale)

    result = runner.invoke(app, ["warm-cache", str(data), "--rebuild"])

    assert result.exit_code == 0, result.output
    loaded = setup_cache.load(data, files)
    assert loaded is not None
    assert loaded.vocab is not None
    assert set(loaded.vocab[0].keys()) == {11, 22}


def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
    runner.invoke(app, ["warm-cache", str(data), "--val-fraction", "0.1"])

    result = runner.invoke(app, ["warm-cache", str(data), "--val-fraction", "0.3"])

    assert result.exit_code == 0, result.output
    assert "vocabulary maps: cache hit" in result.output
    assert "fitting normalizer (streaming)" in result.output
    loaded = setup_cache.load(data, [data])
    assert loaded is not None
    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
    assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers