projectAIpopular/tests/test_llama_server.py

"""T2 llama-server 进程管理单测（封闭：假二进制 + 注入，D11）。"""
import json
import os
import subprocess
import sys
import tempfile
import uuid
from pathlib import Path

import pytest

from runtime.llama_server import LlamaServerManager, LlamaServerError, build_llama_server
from tests._ports import free_port

FAKE_SCRIPT = Path(__file__).parent / "fixtures" / "fake_llama_server.py"
PYTHON = sys.executable


def _make_fake_binary(tmp: Path) -> Path:
    """生成一个 .cmd 包装器：把 venv python + 假脚本当作"二进制"启动。"""
    cmd = tmp / "fake-llama-server.cmd"
    cmd.write_text(
        f'@echo off\r\n"{PYTHON}" "{FAKE_SCRIPT}" %*\r\n',
        encoding="utf-8",
    )
    return cmd


def _make_manager(tmp, binary, port, model, **kw):
    # marker 固定写入系统临时目录；env 仅传文件名（与 fixtures/fake_llama_server.py 对齐）
    marker = Path(tempfile.gettempdir()) / f"fake-llama-marker-{uuid.uuid4().hex}.json"
    env = dict(os.environ)
    env["FAKE_MARKER_NAME"] = marker.name
    return LlamaServerManager(
        binary=str(binary),
        model=str(model),
        port=port,
        hw={"tier": "cpu"},
        health_timeout_s=15.0,
        poll_interval_s=0.2,
        log_dir=str(tmp / "runs"),
        env=env,
        **kw,
    ), marker


def test_build_command_uses_hw_tier():
    m = LlamaServerManager(binary="bin/x.exe", model="models/m.gguf", port=8901,
                           hw={"tier": "gpu12"})
    cmd = m._build_command()
    assert Path(cmd[0]) == Path("bin/x.exe")
    assert cmd[1] == "-m" and Path(cmd[2]) == Path("models/m.gguf")
    assert "--port" in cmd and "8901" in cmd
    assert cmd[cmd.index("-ngl") + 1] == "99"
    assert cmd[cmd.index("-c") + 1] == "32768"


def test_build_command_extra_args_appended():
    m = LlamaServerManager(binary="bin/x.exe", model="models/m.gguf", port=1,
                           hw={"tier": "cpu"}, extra_args=["--cache-reuse", "256"])
    cmd = m._build_command()
    assert cmd[-2:] == ["--cache-reuse", "256"]


def test_start_missing_binary_raises(tmp_path):
    m = LlamaServerManager(binary=str(tmp_path / "nope.exe"), model=str(tmp_path / "m.gguf"),
                           port=free_port())
    with pytest.raises(LlamaServerError):
        m.start()


def test_start_missing_model_raises(tmp_path):
    binary = _make_fake_binary(tmp_path)
    m = LlamaServerManager(binary=str(binary), model=str(tmp_path / "missing.gguf"),
                           port=free_port())
    with pytest.raises(LlamaServerError):
        m.start()


def test_full_lifecycle(tmp_path):
    binary = _make_fake_binary(tmp_path)
    model = tmp_path / "model.gguf"
    model.write_bytes(b"fake")
    port = free_port()
    m, marker = _make_manager(tmp_path, binary, port, model)

    assert m.running is False
    assert m.health() is False  # 无进程时不健康

    ok = m.start()
    assert ok is True
    assert m.running is True
    assert m.health() is True

    # 假脚本确实收到了参数
    data = json.loads(marker.read_text(encoding="utf-8"))
    assert data["port"] == port
    assert data["model"] == str(model)

    # 再 start 幂等（已运行返回健康）
    assert m.start() is True

    m.stop()
    assert m.running is False
    assert m.health() is False


def test_stop_idempotent(tmp_path):
    binary = _make_fake_binary(tmp_path)
    model = tmp_path / "model.gguf"
    model.write_bytes(b"fake")
    m, _ = _make_manager(tmp_path, binary, free_port(), model)
    m.stop()  # 未启动时 stop 不抛
    assert m.running is False


def test_ensure_alive_healthy_no_restart(tmp_path):
    binary = _make_fake_binary(tmp_path)
    model = tmp_path / "model.gguf"
    model.write_bytes(b"fake")
    m, _ = _make_manager(tmp_path, binary, free_port(), model)
    m.start()
    assert m.ensure_alive() is True
    # 不应触发重启
    assert m._restart_count == 0
    m.stop()


def test_ensure_alive_restart_exhausted(tmp_path):
    binary = _make_fake_binary(tmp_path)
    model = tmp_path / "model.gguf"
    model.write_bytes(b"fake")
    m, _ = _make_manager(tmp_path, binary, free_port(), model, max_restarts=0)
    # 未启动：restart 上限 0 -> False
    assert m.ensure_alive() is False


def test_endpoint_format():
    m = LlamaServerManager(binary="x", model="m", port=8901, hw={"tier": "cpu"})
    assert m.endpoint() == "http://127.0.0.1:8901"


def test_build_from_config(tmp_path):
    binary = _make_fake_binary(tmp_path)
    cfg = {
        "binary": str(binary),
        "model": str(tmp_path / "m.gguf"),
        "port": 8999,
        "hw": {"tier": "cpu"},
        "extra_args": ["-fa"],
    }
    m = build_llama_server(cfg)
    assert m.port == 8999
    assert "-fa" in m.extra_args