projectAIpopular/tests/test_proxy_budget.py

"""预算四档渐进干预测试（T-X1）：budget_mode 黄金边界 + 降档选择器。

黄金用例锁死（整数基点 bp，1% = 100bp）：
  projected < 8000bp        -> normal
  8000bp <= projected       -> optimize
  9500bp <= projected       -> cheap
  projected > 10000bp       -> block（try_hold 硬拒兜底）
"""
import pytest

from gateway.model_pool import PoolStore
from gateway.proxy.billing import BUDGET_CHEAP_BP, BUDGET_OPTIMIZE_BP
from gateway.proxy.ledger import Ledger
from gateway.proxy.ledgerutil import _today
from gateway.proxy.routes import _budget_headers, _downgrade_entry

pytest.importorskip("fastapi")

_TS = 1789874000.0   # 固定时间戳（黄金用例确定性）


def _make_ledger(tmp_path, daily_cap_yuan: float = 10.0) -> Ledger:
    led = Ledger.init_db(tmp_path / "budget.sqlite3")
    led.upsert_student("张三", "软件2201", balance_yuan=100.0,
                       daily_cap_yuan=daily_cap_yuan)
    return led


def _set_spent(led: Ledger, student_id: int, spent_milli: int,
               spent_date: str = "") -> None:
    """直接写当日已花（白盒：budget_mode 的输入口径与 try_hold 一致）。"""
    with led._lock, led._connect() as conn:
        conn.execute(
            "UPDATE students SET spent_today_milli = ?, spent_date = ? WHERE id = ?",
            (spent_milli, spent_date, student_id))


def test_threshold_constants_golden():
    """阈值常量黄金锁定：80% / 95%。"""
    assert BUDGET_OPTIMIZE_BP == 8000
    assert BUDGET_CHEAP_BP == 9500


@pytest.mark.parametrize("spent,est,expected", [
    (0, 7999, "normal"),      # 7999bp
    (0, 8000, "optimize"),    # 恰好 8000bp（边界含头）
    (0, 9499, "optimize"),
    (0, 9500, "cheap"),       # 恰好 9500bp（边界含头）
    (0, 10000, "cheap"),      # 恰好达上限 = cheap（不 block）
    (0, 10001, "block"),      # 严格超出才 block
    (5000, 3000, "optimize"), # 8000bp：spent 与 est 合并计算
    (5000, 4500, "cheap"),    # 9500bp
    (5000, 5001, "block"),    # 10001bp
])
def test_budget_mode_golden_boundaries(tmp_path, spent, est, expected):
    led = _make_ledger(tmp_path, daily_cap_yuan=10.0)   # cap = 10000 毫元
    _set_spent(led, 1, spent, spent_date=_today(_TS))   # 记账日 = 判定日（当日口径）
    assert led.budget_mode(1, est, ts=_TS) == expected


def test_budget_mode_zero_cap_means_unlimited(tmp_path):
    led = _make_ledger(tmp_path, daily_cap_yuan=0.0)
    sid = 1
    _set_spent(led, sid, 999999, spent_date=_today(_TS))
    assert led.budget_mode(sid, 10**9, ts=_TS) == "normal"


def test_budget_mode_spent_resets_by_date(tmp_path):
    """spent_date 非今日 -> 当日已花按 0 计（与 try_hold 口径一致）。"""
    led = _make_ledger(tmp_path, daily_cap_yuan=10.0)
    _set_spent(led, 1, 9999, spent_date="2000-01-01")   # 非今日
    assert led.budget_mode(1, 1, ts=1789874000.0) == "normal"


def test_budget_mode_missing_student_normal(tmp_path):
    led = Ledger.init_db(tmp_path / "b.sqlite3")
    assert led.budget_mode(999, 100, ts=1789874000.0) == "normal"


def _make_pool() -> PoolStore:
    import gateway.model_pool as mp
    store = PoolStore()          # 不落盘（path=None 仅内存）
    mp.reset_pool()
    store.upsert({"id": "p1", "name": "旗舰", "tier": "premium", "backend": "openai",
                  "base_url": "https://api.example.com", "model": "big-x",
                  "enabled": True})
    store.upsert({"id": "b1", "name": "实惠", "tier": "budget", "backend": "openai",
                  "base_url": "https://api.example.com/v2", "model": "mid-y",
                  "enabled": True})
    store.upsert({"id": "l1", "name": "本地", "tier": "local",
                  "backend": "llama_server", "base_url": "http://127.0.0.1:8901",
                  "model": "qwen-local", "enabled": True})
    return store


def test_downgrade_optimize_one_tier():
    pool = _make_pool()
    premium = pool.find_by_model("big-x")
    down = _downgrade_entry(pool, premium, "optimize")
    assert down is not None and down["tier"] == "budget"


def test_downgrade_cheap_to_lowest():
    pool = _make_pool()
    premium = pool.find_by_model("big-x")
    down = _downgrade_entry(pool, premium, "cheap")
    assert down is not None and down["tier"] == "local"


def test_downgrade_stops_at_local():
    """已在最低档：cheap/optimize 均不再降（杜绝反向升档）。"""
    pool = _make_pool()
    local = pool.find_by_model("qwen-local")
    assert _downgrade_entry(pool, local, "cheap") is None
    assert _downgrade_entry(pool, local, "optimize") is None


def test_downgrade_skips_disabled_and_mock():
    pool = _make_pool()
    pool.upsert({"id": "b2", "name": "停用", "tier": "budget", "backend": "openai",
                 "base_url": "https://api.example.com/v3", "model": "mid-z",
                 "enabled": False})
    pool.upsert({"id": "mk", "name": "假", "tier": "budget", "backend": "mock",
                 "model": "mock", "enabled": True})
    premium = pool.find_by_model("big-x")
    down = _downgrade_entry(pool, premium, "optimize")
    assert down is not None and down["id"] == "b1"   # 跳过停用与 mock


def test_budget_headers_only_when_abnormal():
    assert _budget_headers("normal") == {}
    assert _budget_headers("") == {}
    assert _budget_headers("optimize") == {"X-Budget-Mode": "optimize"}
    assert _budget_headers("cheap") == {"X-Budget-Mode": "cheap"}