tests/units/test_cli.py

from pathlib import Path
from unittest import mock

import pytest

from mini_transformer import cli

MINIMAL_CFG = """
model:
  name: demo
  best_checkpoint_path: ./checkpoints/best.pt
  latest_checkpoint_path: null
  tokenizer: demo_tok
  d_model: 16
  num_layers: 1
  num_heads: 2
  d_ff: 32
  dropout_rate: 0.0
  vocab_size: 32
  max_seq_len: 32
  pad_id: 0
  bos_id: 1
  eos_id: 2

tokenizer:
  name: demo_tok
  path: ./tokenizer/tokenizer.json
  corpus: null
  dataset: null
  vocab_size: 32
  max_seq_len: 32
  pad_token: "<pad>"
  bos_token: "<bos>"
  eos_token: "<eos>"
  unk_token: "<unk>"
  special_tokens: ["<pad>", "<bos>", "<eos>", "<unk>"]
  pad_id: 0
  bos_id: 1
  eos_id: 2
  unk_id: 3

generation:
  max_new_tokens: 4
  temperature: 1.0
  top_k: null
  top_p: null
  do_sample: false
  presence_penalty: 0.0
  frequency_penalty: 0.0
  no_repeat_ngram: null
  min_steps_before_eos: 0

runtime:
  seed: 42
  device: cpu
  output_dir: ./outputs
  data_dir: ./data
  tokenizer_dir: ./tokenizer
  cache_dir: ./cache
  checkpoint_path: ./checkpoints

input_text: ""
"""


def _make_demo_model(tmp_path: Path) -> Path:
    models_root = tmp_path / "trained_models"
    model_dir = models_root / "demo"
    config_dir = model_dir / "configs"
    config_dir.mkdir(parents=True, exist_ok=True)
    (model_dir / "checkpoints").mkdir(exist_ok=True)
    (model_dir / "tokenizer").mkdir(exist_ok=True)
    (config_dir / "config_inference.yaml").write_text(MINIMAL_CFG)
    return models_root


def test_infer_main_uses_local_model(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys):
    models_root = _make_demo_model(tmp_path)
    monkeypatch.setenv(cli.MODELS_ENV, str(models_root))

    with mock.patch.object(cli, "run_inference", return_value=["hello"]) as run_mock:
        rc = cli.infer_main(["-m", "demo", "-t", "hi there"])

    assert rc == 0
    cfg_passed = run_mock.call_args[0][0]
    assert Path(cfg_passed.model.best_checkpoint_path).is_absolute()
    assert cfg_passed.input_text == "hi there"
    out = capsys.readouterr().out
    assert "hello" in out


def test_infer_main_defaults_to_first_model(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
    models_root = _make_demo_model(tmp_path)
    monkeypatch.setenv(cli.MODELS_ENV, str(models_root))

    with mock.patch.object(cli, "run_inference", return_value=["out"]) as run_mock:
        cli.infer_main(["-t", "text"])
    assert run_mock.called


def test_fetch_main_invokes_snapshot(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
    monkeypatch.setenv(cli.MODELS_ENV, str(tmp_path / "trained_models"))

    with mock.patch("mini_transformer.cli.snapshot_to_local", return_value=tmp_path) as snap:
        rc = cli.fetch_main(
            [
                "AlaBoussoffara/transformer_test",
                "--name",
                "demo",
                "--force",
            ]
        )

    assert rc == 0
    snap.assert_called_once()
    kwargs = snap.call_args.kwargs
    assert kwargs["local_name"] == "demo"
    assert kwargs["force"] is True