Files
AstrAI/tests/test_serve_cli.py
ViperEkura cb21af38ba feat: unify Docker serving configuration in YAML
- Add server.py --config serve.yaml; explicit CLI flags override YAML
- Add scripts/serve.sh and serve_runtime.py for the Compose lifecycle
- Template server/cpu ports and param mounts in docker-compose.yml
- Document schema in docs/developer/docker-serving.md and params guide
- Add tests for runtime parsing and server CLI merge logic
2026-08-21 23:16:45 +08:00

106 lines
3.0 KiB
Python

"""Unit tests for the serving CLI YAML merge logic."""
import click
import pytest
import torch
from click.testing import CliRunner
from scripts.tools.server import (
_merge_yaml_into_kwargs,
_resolve_server_config,
server_command,
)
def _passed() -> dict:
return {
"host": "0.0.0.0",
"port": 8000,
"reload": False,
"param_path": None,
"device": "cuda",
"dtype": "bfloat16",
"max_batch_size": 16,
"max_seq_len": None,
}
def test_yaml_overrides_click_defaults_but_not_explicit_cli(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text(
"server:\n device: cpu\n dtype: float16\n max_batch_size: 8\n",
encoding="utf-8",
)
merged = _merge_yaml_into_kwargs(
str(config_path), _passed(), explicit_keys={"device"}
)
assert merged["device"] == "cuda"
assert merged["dtype"] == "float16"
assert merged["max_batch_size"] == 8
def test_resolve_config_yaml_wins_by_default(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text(
"server:\n port: 9000\n max_seq_len: 2048\n",
encoding="utf-8",
)
resolved = _resolve_server_config(str(config_path), _passed())
assert resolved["port"] == 9000
assert resolved["max_seq_len"] == 2048
assert resolved["device"] == "cuda"
assert resolved["dtype"] == "bfloat16"
def test_resolve_config_rejects_bad_dtype(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text("server:\n dtype: fp8\n", encoding="utf-8")
with pytest.raises(click.UsageError, match="server.dtype"):
_resolve_server_config(str(config_path), _passed())
def test_server_command_rejects_bad_yaml_dtype(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text("server:\n dtype: fp8\n", encoding="utf-8")
result = CliRunner().invoke(server_command, ["--config", str(config_path)])
assert result.exit_code == 2
assert "server.dtype" in result.output
def test_server_command_merges_yaml_and_cli(tmp_path, monkeypatch):
"""Full CLI path: YAML values apply, explicit CLI flags override, args reach run_server."""
config_path = tmp_path / "serve.yaml"
config_path.write_text(
"server:\n device: cpu\n dtype: float16\n max_batch_size: 8\n",
encoding="utf-8",
)
captured = {}
def fake_run_server(**kwargs):
captured.update(kwargs)
monkeypatch.setattr("scripts.tools.server.run_server", fake_run_server)
result = CliRunner().invoke(
server_command,
["--config", str(config_path), "--max_batch_size", "32"],
)
assert result.exit_code == 0, result.output
assert captured["device"] == "cpu"
assert captured["dtype"] == torch.float16
assert captured["max_batch_size"] == 32
assert captured["port"] == 8000
def test_config_option_rejects_missing_file(tmp_path):
result = CliRunner().invoke(
server_command, ["--config", str(tmp_path / "nope.yaml")]
)
assert result.exit_code == 2