projectAIpopular/tests/test_model_pool.py

"""模型池（PoolStore）测试：条目校验/CRUD/角色指派/管线解析/成本分账。"""
import os

import pytest

pytest.importorskip("fastapi")

from fastapi.testclient import TestClient

import gateway.api as ga
import gateway.model_pool as mp
from gateway.model_pool import PoolStore, compute_cost, entry_to_architect_cfg, entry_to_worker_cfg


@pytest.fixture()
def pool(tmp_path):
    """独立文件的全局池（不污染 config/model_pool.json）。"""
    mp.reset_pool()
    store = PoolStore(path=tmp_path / "model_pool.json")
    mp._store = store
    yield store
    mp.reset_pool()
    ga.rebuild_pipeline()


@pytest.fixture()
def client():
    return TestClient(ga.app)


def _entry(**over):
    base = {
        "id": "prem-1", "name": "旗舰模型", "tier": "premium", "backend": "openai",
        "base_url": "https://api.deepseek.com", "model": "deepseek-v4-pro",
        "api_key": os.environ.get("TEST_POOL_KEY", "local-test-only"), "price_in": 1.0, "price_out": 2.0,
        "enabled": True,
    }
    base.update(over)
    return base


# ---------------- PoolStore 单元 ----------------

def test_pool_upsert_and_mask(pool):
    masked = pool.upsert(_entry())
    assert masked["api_key_set"] is True
    assert masked["api_key"] != _entry()["api_key"]    # 明文不打回
    data = pool.list()
    assert data["entries"][0]["model"] == "deepseek-v4-pro"
    assert data["entries"][0]["api_key_set"] is True


def test_pool_upsert_keeps_key_when_blank(pool):
    pool.upsert(_entry())
    pool.upsert(_entry(api_key=""))                    # 前端不回传明文 -> 保留
    assert pool.get("prem-1")["api_key"] == _entry()["api_key"]


def test_pool_validation(pool):
    with pytest.raises(ValueError):
        pool.upsert(_entry(tier="超豪华"))
    with pytest.raises(ValueError):
        pool.upsert(_entry(backend="magic"))
    with pytest.raises(ValueError):
        pool.upsert(_entry(backend="openai", base_url=""))   # 非 mock 缺端点
    with pytest.raises(ValueError):
        pool.upsert(_entry(hack="x"))                        # 未知字段
    with pytest.raises(ValueError):
        pool.upsert(_entry(price_in=-1))


def test_pool_roles_and_resolve(pool):
    pool.upsert(_entry())
    pool.upsert(_entry(id="local-1", tier="local", backend="llama_server",
                       base_url="http://127.0.0.1:8901/v1", model="qwen3.5-4b",
                       price_in=0, price_out=0))
    assert pool.resolve("architect") is None              # 未指派
    pool.set_roles({"architect": "prem-1", "worker": "local-1"})
    assert pool.resolve("architect")["id"] == "prem-1"
    assert pool.resolve("worker")["id"] == "local-1"
    assert pool.resolve("agent") is None
    # 指派不存在的条目
    with pytest.raises(ValueError):
        pool.set_roles({"agent": "ghost"})
    # 删除条目 -> 角色自动清空
    pool.delete("prem-1")
    assert pool.resolve("architect") is None


def test_pool_disabled_entry_not_resolved(pool):
    pool.upsert(_entry(enabled=False))
    pool.set_roles({"architect": "prem-1"})
    assert pool.resolve("architect") is None              # 禁用 -> 回退经典设置


def test_entry_cfg_mapping(pool):
    e = pool.get("prem-1") or _entry()
    acfg = entry_to_architect_cfg(_entry())
    assert acfg["model"] == "deepseek-v4-pro"
    assert acfg["api_key"] == _entry()["api_key"]
    wcfg = entry_to_worker_cfg(_entry())
    assert wcfg["backend"] == "openai"


def test_compute_cost():
    e = {"price_in": 1.0, "price_out": 2.0}
    assert compute_cost(e, 1_000_000, 500_000) == pytest.approx(2.0)
    assert compute_cost({"price_in": 0, "price_out": 0}, 999, 999) == 0.0


# ---------------- API 端点 ----------------

def test_pool_api_crud(pool, client):
    r = client.get("/pool")
    assert r.status_code == 200
    assert r.json()["roles"]["architect"] == ""
    r2 = client.post("/pool", json=_entry())
    assert r2.status_code == 200
    assert len(r2.json()["entries"]) == 1
    # 非法条目 -> 400
    r3 = client.post("/pool", json=_entry(tier="bad"))
    assert r3.status_code == 400
    # 角色指派
    r4 = client.put("/pool/roles", json={"architect": "prem-1"})
    assert r4.status_code == 200
    assert r4.json()["roles"]["architect"] == "prem-1"
    # 删除
    r5 = client.delete("/pool/prem-1")
    assert r5.status_code == 200
    assert r5.json()["roles"]["architect"] == ""


def test_build_pipeline_uses_pool(pool, monkeypatch):
    """池指派应覆盖经典设置，测试 override 最后生效。"""
    pool.upsert(_entry())
    pool.set_roles({"architect": "prem-1"})
    ga.rebuild_pipeline()

    captured = {}

    def fake_build_architect(cfg):
        captured["architect"] = dict(cfg)
        from router_system.architect import ArchitectClient
        return ArchitectClient(model=cfg.get("model", "m"), api_key="k")

    monkeypatch.setattr(ga, "build_architect", fake_build_architect)
    pipe = ga.build_v2_pipeline(worker_cfg_override={"backend": "mock"})
    assert pipe is not None
    assert captured["architect"]["model"] == "deepseek-v4-pro"   # 池条目生效
    ga.rebuild_pipeline()


def test_v2stats_by_model():
    from router_system.v2stats import V2Stats

    class R:
        request_id = "x"
        fast_path = False
        status = "done"
        rounds_used = 1
        api_input_tokens = 1000
        api_output_tokens = 500
        cost_est = 0.002
        model_used = "deepseek-v4-pro"
        route = []

    s = V2Stats()
    s.record(R())
    summary = s.summary()
    bucket = summary["by_model"]["deepseek-v4-pro"]
    assert bucket["requests"] == 1
    assert bucket["input_tokens"] == 1000
    assert bucket["cost_est_usd"] == pytest.approx(0.002)