AstrAI/tests/config/test_cli.py

import dataclasses
import typing as t

import click
import pytest
from click.testing import CliRunner

from astrai.config import TrainConfig, merge_yaml_into_kwargs
from astrai.config.cli import (
    GroupedCommand,
    OptSpec,
    apply_specs,
    option_from_spec,
)


def test_merge_yaml_overrides_defaults_but_not_explicit_cli(tmp_path):
    config_path = tmp_path / "train.yaml"
    config_path.write_text(
        "training:\n"
        "  optimizer: nora_nadamw\n"
        "  max_lr: 2e-4\n"
        "  nora_lr: 0.004\n"
        "  batch_per_device: 8\n",
        encoding="utf-8",
    )
    click_values = {
        "optimizer": "nora_nadamw",
        "max_lr": 3e-4,
        "nora_lr": 5e-3,
        "batch_per_device": 16,
    }

    merged = merge_yaml_into_kwargs(
        str(config_path), click_values, explicit_keys={"batch_per_device"}
    )

    assert merged["max_lr"] == 2e-4
    assert merged["nora_lr"] == 4e-3
    assert merged["batch_per_device"] == 16


def test_merge_yaml_parses_scientific_notation_as_float(tmp_path):
    config_path = tmp_path / "train.yaml"
    config_path.write_text("training:\n  max_lr: 2e-5\n", encoding="utf-8")

    merged = merge_yaml_into_kwargs(
        str(config_path), {"max_lr": 3e-4}, explicit_keys=set()
    )

    assert merged["max_lr"] == 2e-5
    assert isinstance(merged["max_lr"], float)


def test_merge_yaml_ignores_unknown_sections(tmp_path):
    config_path = tmp_path / "train.yaml"
    config_path.write_text("unknown:\n  max_lr: 1.0\n", encoding="utf-8")

    merged = merge_yaml_into_kwargs(str(config_path), {"max_lr": 3e-4})

    assert merged["max_lr"] == 3e-4


def test_type_inference_from_config_fields():
    @click.command()
    @apply_specs(
        [
            OptSpec("n_epoch", "G"),
            OptSpec("rollout_max_policy_lag", "G"),
            OptSpec("pin_memory", "G"),
            OptSpec("metrics", "G"),
        ],
        TrainConfig,
    )
    def cmd(**kwargs):
        pass

    params = {p.name: p for p in cmd.params}
    assert params["n_epoch"].default == 1
    assert params["n_epoch"].type.name == "integer"
    assert params["rollout_max_policy_lag"].default is None
    assert params["rollout_max_policy_lag"].type.name == "integer"
    assert params["pin_memory"].secondary_opts == ["--no-pin_memory"]
    assert params["pin_memory"].default is False  # config default, no override
    assert params["metrics"].multiple
    assert params["metrics"].default == ("loss", "lr", "grad_norm")


def test_spec_overrides_beat_config_defaults():
    @click.command()
    @apply_specs(
        [
            OptSpec("num_workers", "G", default=4),
            OptSpec("parallel_mode", "G", default="fsdp"),
        ],
        TrainConfig,
    )
    def cmd(**kwargs):
        pass

    params = {p.name: p for p in cmd.params}
    assert params["num_workers"].default == 4
    assert params["parallel_mode"].default == "fsdp"


def test_cli_only_flag_pair_and_one_way_flag():
    @click.command()
    @apply_specs(
        [
            OptSpec("muon_nesterov", "G", type=bool, default=True),
            OptSpec("dry_run", "G", is_flag=True, default=False),
        ]
    )
    def cmd(**kwargs):
        pass

    params = {p.name: p for p in cmd.params}
    assert params["muon_nesterov"].secondary_opts == ["--no-muon_nesterov"]
    assert params["muon_nesterov"].default is True
    assert params["dry_run"].is_flag
    assert not params["dry_run"].secondary_opts


def test_uninferrable_type_raises_without_spec_type():
    with pytest.raises(TypeError, match="cannot infer"):
        option_from_spec(OptSpec("mystery", "G"), {})


def test_annotation_styles_old_and_new():
    """Optional[X], Union[X, None], X | None, List[str], and list[str] all
    infer the same click types."""

    @dataclasses.dataclass
    class OldStyle:
        opt_int: t.Optional[int] = None
        opt_float: t.Optional[float] = None
        union_int: t.Union[int, None] = None
        names: t.List[str] = dataclasses.field(default_factory=lambda: ["loss", "lr"])
        flag: bool = True

    @dataclasses.dataclass
    class NewStyle:
        opt_int: int | None = None
        names: list[str] = dataclasses.field(default_factory=lambda: ["loss"])
        flag: bool = False

    @click.command()
    @apply_specs(
        [
            OptSpec("opt_int", "G"),
            OptSpec("opt_float", "G"),
            OptSpec("union_int", "G"),
            OptSpec("names", "G"),
            OptSpec("flag", "G"),
        ],
        OldStyle,
    )
    def old_cmd(**kwargs):
        pass

    params = {p.name: p for p in old_cmd.params}
    assert params["opt_int"].type.name == "integer"
    assert params["opt_int"].default is None
    assert params["opt_float"].type.name == "float"
    assert params["union_int"].type.name == "integer"
    assert params["names"].multiple
    assert params["names"].type.name == "text"
    assert params["flag"].secondary_opts == ["--no-flag"]
    assert params["flag"].default is True

    @click.command()
    @apply_specs(
        [
            OptSpec("opt_int", "G"),
            OptSpec("names", "G"),
            OptSpec("flag", "G"),
        ],
        NewStyle,
    )
    def new_cmd(**kwargs):
        pass

    params = {p.name: p for p in new_cmd.params}
    assert params["opt_int"].type.name == "integer"
    assert params["names"].multiple
    assert params["names"].default == ("loss",)
    assert params["flag"].secondary_opts == ["--no-flag"]
    assert params["flag"].default is False


def test_stringified_annotations_resolve_via_get_type_hints():
    """PEP 563 modules (``from __future__ import annotations``) leave
    ``Field.type`` as a string; hints resolution still infers types."""

    @dataclasses.dataclass
    class FutureStyle:
        opt_int: "t.Optional[int]" = None
        count: "int" = 3

    @click.command()
    @apply_specs(
        [
            OptSpec("opt_int", "G"),
            OptSpec("count", "G"),
        ],
        FutureStyle,
    )
    def cmd(**kwargs):
        pass

    params = {p.name: p for p in cmd.params}
    assert params["opt_int"].type.name == "integer"
    assert params["count"].type.name == "integer"
    assert params["count"].default == 3


def test_help_order_follows_spec_table():
    @click.command(cls=GroupedCommand)
    @apply_specs(
        [
            OptSpec("first", "G", type=int, default=1),
            OptSpec("second", "G", type=int, default=2),
        ]
    )
    def cmd(**kwargs):
        pass

    result = CliRunner().invoke(cmd, ["--help"])
    assert result.exit_code == 0
    assert result.output.index("--first") < result.output.index("--second")