onf-c2/tests/test_cli.py

"""Tests for onf.graph.cli — python -m onf.graph's argument parsing, --dry-run
contract, suite validation, exit codes, and full end-to-end build/train runs through the CLI
against the fabricated tiny HDF5 fixture from tests.test_builders (not real LIBERO data).

Only build and train are supported subcommands.
"""

from __future__ import annotations

import json
import subprocess
import sys
from pathlib import Path

import pytest
from test_builders import _write_fixture

from onf.graph.core import schema
from onf.graph.cli import (
    CLIError,
    _resolve_suites,
    build_parser,
    load_configured_suites,
    main,
)
from onf.graph.core.edges import EdgeSet
from onf.graph.core.nodes import NodeTable


# ==================================================================================================
# isolate ONF_OUTPUTS -- every test in this file must never touch the real repo outputs/ tree
# (tests/test_builders.py's own convention, followed here verbatim).
# ==================================================================================================
@pytest.fixture(autouse=True)
def _isolate_outputs(tmp_path, monkeypatch):
    monkeypatch.setenv("ONF_OUTPUTS", str(tmp_path / "outputs"))
    return tmp_path / "outputs"


@pytest.fixture
def outputs_root(tmp_path):
    return tmp_path / "outputs"


# ==================================================================================================
# suites.yaml / suite resolution
# ==================================================================================================
def test_load_configured_suites_matches_yaml():
    suites = load_configured_suites()
    assert suites == ["object", "spatial", "goal", "long"]


def test_resolve_suites_all_expands():
    configured = ["object", "spatial", "goal", "long"]
    assert _resolve_suites("all", configured) == configured


def test_resolve_suites_single_name():
    configured = ["object", "spatial", "goal", "long"]
    assert _resolve_suites("spatial", configured) == ["spatial"]


def test_resolve_suites_unknown_raises_cli_error():
    with pytest.raises(CLIError, match="unknown suite"):
        _resolve_suites("bogus", ["object", "spatial", "goal", "long"])


# ==================================================================================================
# argparse construction -- every subcommand parses its own documented flags
# ==================================================================================================
def test_build_parser_has_every_subcommand():
    parser = build_parser()
    sub_actions = [a for a in parser._subparsers._group_actions if hasattr(a, "choices")]
    names = set(sub_actions[0].choices.keys())
    assert names == {"build", "train"}


@pytest.mark.parametrize("command", ["build", "train"])
def test_subcommand_help_exits_zero_and_is_non_trivial(capsys, command):
    parser = build_parser()
    with pytest.raises(SystemExit) as exc:
        parser.parse_args([command, "--help"])
    assert exc.value.code == 0
    out = capsys.readouterr().out
    assert len(out) > 200      # a real, substantive description -- not a one-liner


def test_build_args_parse_all_flags():
    parser = build_parser()
    args = parser.parse_args([
        "build", "--suite", "long", "--coarsen", "7", "--device", "cuda:0",
        "--limit-demos", "3", "--limit-tasks", "2", "--hdf5-dir", "/tmp/x",
        "--out", "/tmp/o", "--name", "n", "--seed", "5", "--quiet", "--dry-run", "--force",
    ])
    assert args.suite == "long" and args.coarsen == 7 and args.device == "cuda:0"
    assert args.limit_demos == 3 and args.limit_tasks == 2 and args.hdf5_dir == "/tmp/x"
    assert args.out == "/tmp/o" and args.name == "n" and args.seed == 5
    assert args.quiet is True and args.dry_run is True and args.force is True
    assert args.resume is None


def test_missing_required_suite_is_argparse_usage_error():
    parser = build_parser()
    with pytest.raises(SystemExit) as exc:
        parser.parse_args(["build"])
    assert exc.value.code == 2      # argparse's own usage-error convention


# ==================================================================================================
# --suite all expansion (through main(), not just _resolve_suites in isolation)
# ==================================================================================================
def test_suite_all_expands_in_dry_run(capsys, outputs_root):
    rc = main(["build", "--suite", "all", "--dry-run"])
    assert rc == 0
    out = capsys.readouterr().out
    for suite in ("object", "spatial", "goal", "long"):
        assert f"build[{suite}]" in out
    assert not outputs_root.exists()


# ==================================================================================================
# unknown suite -> clear message, exit code 1
# ==================================================================================================
def test_unknown_suite_fails_with_exit_1(capsys):
    rc = main(["build", "--suite", "not_a_real_suite", "--dry-run"])
    assert rc == 1
    err = capsys.readouterr().err
    assert "unknown suite" in err
    assert "not_a_real_suite" in err


# ==================================================================================================
# --dry-run: resolved config printed, outputs/ never touched -- every subcommand
# ==================================================================================================
@pytest.mark.parametrize("argv", [
    ["build", "--suite", "object"],
    ["train", "--suite", "object"],
])
def test_dry_run_touches_nothing_and_prints_config(capsys, outputs_root, argv):
    rc = main([*argv, "--dry-run"])
    assert rc == 0
    out = capsys.readouterr().out
    assert "fully-resolved config" in out
    assert not outputs_root.exists()


def test_dry_run_build_config_has_every_default_materialized(capsys, outputs_root):
    rc = main(["build", "--suite", "object", "--dry-run"])
    assert rc == 0
    out = capsys.readouterr().out
    start = out.index("{")
    end = out.index("planned stages")
    config = json.loads(out[start:end])
    # every GraphConfig field must be present -- "every default materialized", not just the ones the
    # user happened to pass on the command line
    for field in ("coarsen", "k_sibling", "k_align", "hidden", "layers", "agg",
                  "readout", "temp", "topm"):
        assert field in config, field
    assert config["suite"] == "object"
    assert config["hdf5_dir"]      # resolved, non-empty, never a bare placeholder


# ==================================================================================================
# real, tiny, end-to-end build through the CLI (fabricated fixture, no real LIBERO data)
# ==================================================================================================
def test_build_end_to_end_leaves_well_formed_run_dir(tmp_path, outputs_root, capsys):
    fixture = tmp_path / "hdf5"
    _write_fixture(fixture, n_tasks=2, n_demos=3, base_t=40)

    rc = main(["build", "--suite", "object", "--hdf5-dir", str(fixture), "--coarsen", "3"])
    assert rc == 0

    run_dirs = sorted((outputs_root / "object").glob("graph_*"))
    assert len(run_dirs) == 1, run_dirs
    run_dir = run_dirs[0]

    for fname in ("config.json", "manifest.json", "metrics.json", "log.txt"):
        assert (run_dir / fname).exists(), fname
    manifest = json.loads((run_dir / "manifest.json").read_text())
    assert manifest["status"] == "complete"

    nodes_path = run_dir / "artifacts" / schema.NODES_NPZ
    edges_path = run_dir / "artifacts" / schema.EDGES_NPZ
    assert nodes_path.exists() and edges_path.exists()
    table = NodeTable.load(nodes_path)
    table.validate()
    es = EdgeSet.load(edges_path)
    es.validate(table)
    assert table.n_demos == 6 and table.n_tasks == 2

    # latest was repointed at this build (build's StageLogger uses update_latest=True)
    assert (outputs_root / "object" / "latest").resolve() == run_dir.resolve()

    out = capsys.readouterr().out
    assert "built" in out


def test_build_skips_identical_config_unless_forced(tmp_path, outputs_root, capsys):
    fixture = tmp_path / "hdf5"
    _write_fixture(fixture, n_tasks=1, n_demos=2, base_t=40)
    argv = ["build", "--suite", "object", "--hdf5-dir", str(fixture), "--coarsen", "3"]

    assert main(argv) == 0
    capsys.readouterr()
    run_dirs_after_first = sorted((outputs_root / "object").glob("graph_*"))
    assert len(run_dirs_after_first) == 1

    assert main(argv) == 0     # identical config -> SKIP, no second run dir
    out = capsys.readouterr().out
    assert "SKIP" in out
    run_dirs_after_second = sorted((outputs_root / "object").glob("graph_*"))
    assert len(run_dirs_after_second) == 1

    # --force -> a genuine rebuild, not a skip (StageLogger's run-dir timestamp is second-resolution
    # -- see onf.graph.report.TIMESTAMP_FMT's own docstring note -- so back-to-back calls in a fast
    # test CAN collide on the same directory name; what --force actually guarantees is "did not skip",
    # which the printed status -- not the directory count -- is the reliable way to check here).
    assert main([*argv, "--force"]) == 0
    out = capsys.readouterr().out
    assert "SKIP" not in out
    assert "built" in out


def test_build_suite_all_builds_every_configured_suite_against_same_fixture(tmp_path, outputs_root):
    """--suite all combined with --hdf5-dir points every configured suite at the SAME tiny
    fixture -- a cheap way to exercise the multi-suite loop end-to-end without four real datasets."""
    fixture = tmp_path / "hdf5"
    _write_fixture(fixture, n_tasks=1, n_demos=2, base_t=40)

    rc = main(["build", "--suite", "all", "--hdf5-dir", str(fixture), "--coarsen", "3"])
    assert rc == 0
    for suite in ("object", "spatial", "goal", "long"):
        run_dirs = sorted((outputs_root / suite).glob("graph_*"))
        assert len(run_dirs) == 1, (suite, run_dirs)


# ==================================================================================================
# train: requires a build first
# ==================================================================================================
def test_train_requires_graph_built_first(outputs_root):
    rc = main(["train", "--suite", "object", "--epochs", "1"])
    assert rc == 1


def test_train_end_to_end_writes_checkpoint_next_to_graph(tmp_path, outputs_root):
    fixture = tmp_path / "hdf5"
    _write_fixture(fixture, n_tasks=2, n_demos=3, base_t=40)
    assert main(["build", "--suite", "object", "--hdf5-dir", str(fixture), "--coarsen", "3"]) == 0

    graph_run_dir = sorted((outputs_root / "object").glob("graph_*"))[0]
    rc = main(["train", "--suite", "object", "--epochs", "1",
               "--n-query", "16", "--n-eval", "4"])
    assert rc == 0

    # checkpoint lands in the BUILD's artifacts dir, not the train run's own -- update_latest=False
    # on train's StageLogger must leave latest pointed at the build.
    checkpoint = graph_run_dir / "artifacts" / schema.HEAD_NPZ
    assert checkpoint.exists()
    assert (outputs_root / "object" / "latest").resolve() == graph_run_dir.resolve()

    train_run_dirs = sorted((outputs_root / "object").glob("train_*"))
    assert len(train_run_dirs) == 1


# ==================================================================================================
# __main__.py -- the actual python -m onf.graph entry point (subprocess, light coverage)
# ==================================================================================================
def test_python_dash_m_onf_graph_help(monkeypatch):
    env = dict(**{"PYTHONPATH": "src"})
    import os

    full_env = dict(os.environ)
    full_env.update(env)
    full_env.setdefault("OMP_NUM_THREADS", "4")
    full_env.setdefault("MKL_NUM_THREADS", "4")
    full_env.setdefault("OPENBLAS_NUM_THREADS", "4")
    result = subprocess.run(
        [sys.executable, "-m", "onf.graph", "--help"],
        cwd=str(Path(__file__).resolve().parents[1]), env=full_env,
        capture_output=True, text=True, timeout=30,
    )
    assert result.returncode == 0
    assert "build" in result.stdout and "train" in result.stdout


# ==================================================================================================
# small helper-level unit tests (no I/O)
# ==================================================================================================
def test_pick_prefers_explicit_then_resumed_then_default():
    from onf.graph.cli import _pick

    assert _pick(5, {"x": 1}, "x", 0) == 5
    assert _pick(None, {"x": 1}, "x", 0) == 1
    assert _pick(None, {}, "x", 0) == 0


def test_train_resolved_config_keeps_the_resolved_graph_dir(capsys, outputs_root):
    """GraphConfig carries its own graph_dir field (onf.config's GR_GRAPH_DIR override), which
    collides with the key train resolves. Splatting the dataclass last used to overwrite the
    resolved path with GraphConfig's default "", so config.json recorded an empty graph_dir and the
    --force/skip guard compared runs on a key that was always blank."""
    rc = main(["train", "--suite", "object", "--graph-dir", "/tmp/some_graph", "--dry-run"])
    assert rc == 0
    out = capsys.readouterr().out
    config = json.loads(out[out.index("{"):out.index("planned stages")])
    assert config["graph_dir"] == "/tmp/some_graph"


def test_fmt_metric_shows_ci_spread():
    from onf.graph.cli import _fmt_metric

    s = _fmt_metric({"lo": 0.1, "point": 0.2, "hi": 0.3})
    assert "0.2" in s and "0.1" in s and "0.3" in s