tests/unit/test_cli.py

"""CLI smoke tests with mocked LoRA fitness."""

from __future__ import annotations

import json
from pathlib import Path
from unittest.mock import patch

from ahdcma.cli import run_task


def _fake_fitness(x, task, *, num_steps, seed):
    # Return -accuracy with a smooth dependence on x so the optimizer
    # has signal but the test runs in milliseconds.
    return float(-(0.5 - 0.4 * (x[0] - 0.5) ** 2))


def test_run_lora_task_with_mock(tmp_path: Path) -> None:
    with patch("ahdcma.cli.run_task.lora_fitness", side_effect=_fake_fitness):
        summary = run_task.run_lora_task(
            "ahdcma",
            "cifar100_vit",
            seed=0,
            pop=4,
            gens=2,
            num_steps=1,
            output_dir=tmp_path,
        )
    assert summary["best_f"] < 0  # accuracy is positive
    assert summary["best_accuracy"] > 0
    assert summary["n_generations"] == 3  # init + 2 generations
    out_path = tmp_path / summary["run_id"]
    assert (out_path / "result.json").exists()
    assert (out_path / "evaluations.json").exists()
    payload = json.loads((out_path / "result.json").read_text())
    assert payload["algo"] == "ahdcma"
    assert payload["task"] == "cifar100_vit"


def test_run_lora_task_unknown_algo_rejected(tmp_path: Path) -> None:
    import pytest

    with pytest.raises(KeyError):
        run_task.run_lora_task(
            "not_a_real_algo",
            "cifar100_vit",
            seed=0,
            pop=2,
            gens=1,
            num_steps=1,
            output_dir=tmp_path,
        )