giant/tests/test_steps_to_parquet_parallel.py

import json
from pathlib import Path

from scripts import steps_to_parquet_parallel

run_parallel = steps_to_parquet_parallel.run_parallel
resolve_destination = steps_to_parquet_parallel.resolve_destination
latest_schema_tag = steps_to_parquet_parallel.latest_schema_tag
DestinationError = steps_to_parquet_parallel.DestinationError


def _write_fake_executable(tmp_path: Path, marker_dir: Path) -> Path:
    """A stand-in for steps_to_parquet.py: records its own start/end time per
    input file (so tests can check overlap), then exits 1 if "fail" is in the
    filename, else 0. Accepts (and ignores) the same flags as the real script.
    """
    fake = tmp_path / "fake_steps_to_parquet.py"
    fake.write_text(
        f"""
import json, sys, time
from pathlib import Path

marker_dir = Path({str(marker_dir)!r})
root_file = sys.argv[1]
name = Path(root_file).stem
start = time.time()
time.sleep(0.2)
end = time.time()
(marker_dir / f"{{name}}.json").write_text(json.dumps({{"start": start, "end": end}}))
print(f"converted {{root_file}}")
sys.exit(1 if "fail" in name else 0)
"""
    )
    return fake


def test_runs_one_job_per_file_and_reports_success(tmp_path):
    marker_dir = tmp_path / "markers"
    marker_dir.mkdir()
    fake = _write_fake_executable(tmp_path, marker_dir)

    files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(3)]
    results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)

    assert {r[0] for r in results} == set(files)
    assert all(code == 0 for _, code, _, _ in results)
    assert all((marker_dir / f"{Path(f).stem}.json").exists() for f in files)


def test_failures_are_reported_with_nonzero_exit_code(tmp_path):
    marker_dir = tmp_path / "markers"
    marker_dir.mkdir()
    fake = _write_fake_executable(tmp_path, marker_dir)

    files = [str(tmp_path / "shard-000.root"), str(tmp_path / "shard-fail.root")]
    results = run_parallel(files, jobs=4, steps_to_parquet_path=fake)

    codes = {Path(f).stem: code for f, code, _, _ in results}
    assert codes["shard-000"] == 0
    assert codes["shard-fail"] == 1


def test_jobs_caps_concurrency(tmp_path):
    marker_dir = tmp_path / "markers"
    marker_dir.mkdir()
    fake = _write_fake_executable(tmp_path, marker_dir)

    files = [str(tmp_path / f"shard-{i:03d}.root") for i in range(6)]
    run_parallel(files, jobs=2, steps_to_parquet_path=fake)

    intervals = []
    for f in files:
        data = json.loads((marker_dir / f"{Path(f).stem}.json").read_text())
        intervals.append((data["start"], data["end"]))

    # max number of intervals overlapping any given instant must not exceed jobs
    events = sorted([(s, 1) for s, _ in intervals] + [(e, -1) for _, e in intervals])
    concurrent = 0
    peak = 0
    for _, delta in events:
        concurrent += delta
        peak = max(peak, concurrent)
    assert peak <= 2


def test_default_jobs_is_four():
    args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
    assert args.jobs == 4


def test_output_for_is_passed_through_as_output_flag(tmp_path):
    marker_dir = tmp_path / "markers"
    marker_dir.mkdir()
    fake = tmp_path / "fake_steps_to_parquet.py"
    fake.write_text(
        f"""
import sys
from pathlib import Path

marker_dir = Path({str(marker_dir)!r})
root_file = sys.argv[1]
name = Path(root_file).stem
output = sys.argv[sys.argv.index("--output") + 1] if "--output" in sys.argv else ""
(marker_dir / f"{{name}}.txt").write_text(output)
sys.exit(0)
"""
    )
    root_file = str(tmp_path / "shard-000.root")
    dest = tmp_path / "processed" / "shard-000.parquet"
    run_parallel(
        [root_file], jobs=1, steps_to_parquet_path=fake, output_for={root_file: dest}
    )
    assert (marker_dir / "shard-000.txt").read_text() == str(dest)


def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
    root_file = tmp_path / "raw" / "steps" / "gen1" / "pbwo4" / "shard-000.root"
    root_file.parent.mkdir(parents=True)
    root_file.touch()
    for schema in schemas or []:
        (tmp_path / "processed" / "steps" / "gen1" / schema).mkdir(parents=True)
    return root_file


def test_resolve_destination_uses_latest_schema(tmp_path):
    root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
    dest = resolve_destination(root_file, tmp_path, schema_override=None)
    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"


def test_resolve_destination_schema_override_wins(tmp_path):
    root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
    dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"


def test_resolve_destination_errors_without_any_schema(tmp_path):
    root_file = _make_dataset(tmp_path, schemas=[])
    try:
        resolve_destination(root_file, tmp_path, schema_override=None)
        assert False, "expected DestinationError"
    except DestinationError:
        pass


def test_resolve_destination_errors_outside_dataset_root(tmp_path):
    other_root = tmp_path / "other"
    other_root.mkdir()
    stray_file = other_root / "shard-000.root"
    stray_file.touch()
    dataset_root = tmp_path / "dataset"
    dataset_root.mkdir()
    try:
        resolve_destination(stray_file, dataset_root, schema_override=None)
        assert False, "expected DestinationError"
    except DestinationError:
        pass


def test_resolve_destination_errors_on_wrong_shape(tmp_path):
    # missing the <gen> segment
    bad_file = tmp_path / "raw" / "steps" / "pbwo4" / "shard-000.root"
    bad_file.parent.mkdir(parents=True)
    bad_file.touch()
    try:
        resolve_destination(bad_file, tmp_path, schema_override=None)
        assert False, "expected DestinationError"
    except DestinationError:
        pass


def test_latest_schema_tag_returns_none_when_missing(tmp_path):
    assert latest_schema_tag(tmp_path / "does" / "not" / "exist") is None


def test_dataset_root_and_schema_flags_default(tmp_path):
    args = steps_to_parquet_parallel.build_parser().parse_args(["dummy.root"])
    assert args.dataset_root == "/ceph/lbogner/geant_steps"
    assert args.schema is None