tests/units/test_server.py

from pathlib import Path

from fastapi.testclient import TestClient

from mini_transformer.apps import server
from mini_transformer.model_loader import MODEL_NAME_ENV, MODELS_ENV

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_generate_endpoint_uses_local_model(monkeypatch, tmp_path):
    models_root = _make_demo_model(tmp_path)
    monkeypatch.setenv(MODELS_ENV, str(models_root))
    monkeypatch.setenv(MODEL_NAME_ENV, "demo")
    monkeypatch.setattr(server, "run_inference", lambda cfg: ["ok"])

    client = TestClient(server.app)
    response = client.get("/generate", params={"text": "hello"})
    assert response.status_code == 200
    assert response.json()["outputs"] == ["ok"]