tests/test_scaling.py

from __future__ import annotations

from pathlib import Path

from dovla_cil.experiments.scaling import (
    ScalingExperiment,
    parse_k_values,
    read_scaling_csv,
    run_scaling_experiment,
)
from dovla_cil.utils.io import read_json


def test_parse_k_values() -> None:
    assert parse_k_values("1,2,4") == (1, 2, 4)


def test_tiny_scaling_run_writes_csv_and_plots(tmp_path: Path) -> None:
    summary = run_scaling_experiment(
        ScalingExperiment(
            backend="toy",
            tasks="builtins",
            output_dir=tmp_path,
            total_records=4,
            k_values=(1, 2),
            epochs=1,
            seed=0,
            shard_size=8,
            batch_groups=2,
            records_per_group=2,
            hidden_dim=32,
            eval_num_tasks=2,
            device="auto",
        )
    )

    rows = read_scaling_csv(tmp_path / "scaling_results.csv")
    regression = read_json(tmp_path / "scaling_regression.json")

    assert len(rows) == 2
    assert [row["k"] for row in rows] == ["1", "2"]
    assert Path(summary["aggregate_csv"]).exists()
    assert "ranking_acc" in regression
    for filename in (
        "success_rate_vs_k.png",
        "ranking_acc_vs_k.png",
        "instruction_switch_acc_vs_k.png",
        "effect_mae_vs_k.png",
        "regret_ece_vs_k.png",
    ):
        assert (tmp_path / filename).exists()
        assert (tmp_path / filename).stat().st_size > 0
    assert (tmp_path / "k_0001" / "metrics.json").exists()
    assert (tmp_path / "k_0002" / "metrics.json").exists()