200 lines
9.2 KiB
Python
200 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
||
"""llm 额度治理离线单测：全程打桩 _now/_read_window/_bump_window/chat/time.sleep，
|
||
不碰真 DB、不发真请求、不做任何真实 LLM 调用。直接 `.venv/bin/python test_quota.py` 跑（无需 pytest）。
|
||
|
||
覆盖：window_key 归窗 / seconds_to_next_window 边界 / cost_usd 计价(含缓存折扣) /
|
||
chat_governed 六条路由（首选成功、$24 预算耗尽切链、敏感换模型、不可用换模型、全链失败、6000 次达上限睡窗）。"""
|
||
import pathlib
|
||
import sys
|
||
import types
|
||
from datetime import datetime, timedelta
|
||
|
||
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
|
||
import llm # noqa: E402
|
||
|
||
# 费率缓存预置为兜底表：cost_usd/chat_governed 记账时 get_pricing() 直接命中缓存，绝不触网
|
||
# （硬约束：单测不发真请求；兜底费率与 New-API 快照一致，计价结果确定可断言）
|
||
llm._PRICING_CACHE = dict(llm.PRICING_FALLBACK)
|
||
|
||
# 成功调用的假 usage（M3 计价：(100)+50×4=300 单位 → 成本>0，用于断言 MiniMax 计入预算）
|
||
_USAGE = {"prompt_tokens": 100, "completion_tokens": 50,
|
||
"prompt_tokens_details": {"cached_tokens": 0}}
|
||
|
||
|
||
def _install(now, reads, chat, sleep=None):
|
||
"""把一组桩装到 llm 模块全局，返回 bumps 记录列表（每项 = (window_key, add_usd)）。
|
||
now：datetime（固定）或无参可调用（动态时钟）；reads：{window_key:(usd,calls)} 或可调用。"""
|
||
bumps = []
|
||
llm._now = now if callable(now) else (lambda: now)
|
||
llm._read_window = (lambda wk: reads(wk)) if callable(reads) \
|
||
else (lambda wk: reads.get(wk, (0.0, 0)))
|
||
|
||
def fake_bump(wk, add):
|
||
bumps.append((wk, add))
|
||
return (add, 1) # 返回值 chat_governed 不消费，占位即可
|
||
llm._bump_window = fake_bump
|
||
llm.chat = chat
|
||
# 只替换 time.sleep（chat 被打桩后 chat_governed 路径不再用 time.time），不污染真 time 模块
|
||
llm.time = types.SimpleNamespace(sleep=(sleep if sleep else (lambda s: None)))
|
||
return bumps
|
||
|
||
|
||
def _ok_chat(prompt, model=None, **kw):
|
||
"""万能成功桩：任何模型都返回内容。"""
|
||
return "ok", dict(_USAGE)
|
||
|
||
|
||
def test_window_key():
|
||
d = lambda h: datetime(2026, 7, 16, h, 30)
|
||
assert llm.window_key(d(0)) == "2026-07-16T00", "0点应归 T00"
|
||
assert llm.window_key(d(4)) == "2026-07-16T00", "4点应归 T00"
|
||
assert llm.window_key(d(5)) == "2026-07-16T05", "5点应归 T05"
|
||
assert llm.window_key(d(9)) == "2026-07-16T05", "9点应归 T05"
|
||
assert llm.window_key(d(10)) == "2026-07-16T10", "10点应归 T10"
|
||
assert llm.window_key(d(14)) == "2026-07-16T10", "14点应归 T10"
|
||
assert llm.window_key(d(15)) == "2026-07-16T15", "15点应归 T15"
|
||
assert llm.window_key(d(19)) == "2026-07-16T15", "19点应归 T15"
|
||
assert llm.window_key(d(20)) == "2026-07-16T20", "20点应归 T20"
|
||
assert llm.window_key(d(23)) == "2026-07-16T20", "23点应归 T20"
|
||
|
||
|
||
def test_seconds_to_next_window():
|
||
def near(got, want):
|
||
assert abs(got - want) <= 1, f"距下一窗秒数应≈{want}，实得{got}"
|
||
near(llm.seconds_to_next_window(datetime(2026, 7, 16, 3, 0)), 7200) # 03:00→05:00
|
||
near(llm.seconds_to_next_window(datetime(2026, 7, 16, 19, 30)), 1800) # 19:30→20:00
|
||
near(llm.seconds_to_next_window(datetime(2026, 7, 16, 21, 0)), 10800) # 21:00→次日00:00
|
||
near(llm.seconds_to_next_window(datetime(2026, 7, 16, 23, 59)), 60) # 23:59→次日00:00
|
||
|
||
|
||
def test_cost_usd():
|
||
# M3 无缓存：billable=(1000)+500×4=3000 → 0.15×3000/500000=0.0009
|
||
c1 = llm.cost_usd("MiniMax-M3", {"prompt_tokens": 1000, "completion_tokens": 500,
|
||
"prompt_tokens_details": {"cached_tokens": 0}})
|
||
assert abs(c1 - 0.0009) < 1e-9, f"M3 无缓存成本应=0.0009，实得{c1}"
|
||
# M3 有缓存 400：billable=(1000-400)+400×0.2+500×4=2680 → 0.15×2680/500000=0.000804
|
||
c2 = llm.cost_usd("MiniMax-M3", {"prompt_tokens": 1000, "completion_tokens": 500,
|
||
"prompt_tokens_details": {"cached_tokens": 400}})
|
||
assert abs(c2 - 0.000804) < 1e-9, f"M3 缓存折扣后成本应=0.000804，实得{c2}"
|
||
|
||
|
||
def test_current_quota_policy():
|
||
"""额度策略固定为每 5 小时 $24/6000，防止实现和测试再次漂移。"""
|
||
assert llm.WINDOW_BUDGET_USD == 24.0, "当前 MiniMax 每窗预算必须是 $24"
|
||
assert llm.WINDOW_CALL_CAP == 6000, "当前每窗总调用上限必须是 6000"
|
||
|
||
|
||
def test_route_first_success():
|
||
"""①usd=0/calls=0：链首 M3 成功 → 返回 MiniMax-M3、记账 usd>0、bump 一次到本窗。"""
|
||
calls = []
|
||
|
||
def chat(prompt, model=None, **kw):
|
||
calls.append(model)
|
||
return "ok", dict(_USAGE)
|
||
bumps = _install(datetime(2026, 7, 16, 12, 0), {"2026-07-16T10": (0.0, 0)}, chat)
|
||
content, usage, used = llm.chat_governed("p")
|
||
assert used == "MiniMax-M3", f"首选应为 M3，实得{used}"
|
||
assert content == "ok" and calls[0] == "MiniMax-M3", "第一个尝试的模型应是 M3"
|
||
assert len(bumps) == 1 and bumps[0][0] == "2026-07-16T10", "应向本窗 T10 记账一次"
|
||
assert bumps[0][1] > 0, f"M3 是 MiniMax，记账成本应>0，实得{bumps[0][1]}"
|
||
|
||
|
||
def test_route_budget_exhausted():
|
||
"""②usd 超过 $24 预算：链首变 glm-5.2（摘掉 MiniMax）、记账增量=0。"""
|
||
calls = []
|
||
|
||
def chat(prompt, model=None, **kw):
|
||
calls.append(model)
|
||
return "ok", dict(_USAGE)
|
||
bumps = _install(
|
||
datetime(2026, 7, 16, 12, 0),
|
||
{"2026-07-16T10": (llm.WINDOW_BUDGET_USD + 0.5, 0)},
|
||
chat,
|
||
)
|
||
content, usage, used = llm.chat_governed("p")
|
||
assert calls[0] == "glm-5.2", f"预算耗尽后链首应为 glm-5.2，实得{calls[0]}"
|
||
assert used == "glm-5.2", f"应由 glm-5.2 产出，实得{used}"
|
||
assert bumps[0][1] == 0.0, f"glm 非 MiniMax，记账增量应=0，实得{bumps[0][1]}"
|
||
|
||
|
||
def test_route_sensitive_then_next():
|
||
"""③M3 抛 SensitiveError、M2.7 成功 → 返回 MiniMax-M2.7。"""
|
||
calls = []
|
||
|
||
def chat(prompt, model=None, **kw):
|
||
calls.append(model)
|
||
if model == "MiniMax-M3":
|
||
raise llm.SensitiveError("敏感")
|
||
return "ok", dict(_USAGE)
|
||
bumps = _install(datetime(2026, 7, 16, 12, 0), {"2026-07-16T10": (0.0, 0)}, chat)
|
||
content, usage, used = llm.chat_governed("p")
|
||
assert calls == ["MiniMax-M3", "MiniMax-M2.7"], f"应 M3 敏感后换 M2.7，实得{calls}"
|
||
assert used == "MiniMax-M2.7", f"应由 M2.7 产出，实得{used}"
|
||
assert bumps[0][1] > 0, "M2.7 是 MiniMax，记账成本应>0"
|
||
|
||
|
||
def test_route_runtime_then_next():
|
||
"""④glm 抛 RuntimeError、deepseek 成功 → 返回 deepseek-v4-flash（前两 MiniMax 先敏感耗尽）。"""
|
||
calls = []
|
||
|
||
def chat(prompt, model=None, **kw):
|
||
calls.append(model)
|
||
if model in ("MiniMax-M3", "MiniMax-M2.7"):
|
||
raise llm.SensitiveError("敏感")
|
||
if model == "glm-5.2":
|
||
raise RuntimeError("不可用")
|
||
return "ok", dict(_USAGE)
|
||
bumps = _install(datetime(2026, 7, 16, 12, 0), {"2026-07-16T10": (0.0, 0)}, chat)
|
||
content, usage, used = llm.chat_governed("p")
|
||
assert calls == ["MiniMax-M3", "MiniMax-M2.7", "glm-5.2", "deepseek-v4-flash"], \
|
||
f"应走完全链到 deepseek，实得{calls}"
|
||
assert used == "deepseek-v4-flash", f"应由 deepseek 产出，实得{used}"
|
||
assert bumps[0][1] == 0.0, "deepseek 非 MiniMax，记账增量应=0"
|
||
|
||
|
||
def test_route_all_fail():
|
||
"""⑤全链抛 SensitiveError → 返回 (None,None,None)、绝不记账。"""
|
||
def chat(prompt, model=None, **kw):
|
||
raise llm.SensitiveError("全敏感")
|
||
bumps = _install(datetime(2026, 7, 16, 12, 0), {"2026-07-16T10": (0.0, 0)}, chat)
|
||
content, usage, used = llm.chat_governed("p")
|
||
assert (content, usage, used) == (None, None, None), "全链失败应返回三元 None"
|
||
assert bumps == [], "全链失败不得记账"
|
||
|
||
|
||
def test_route_call_cap_sleep():
|
||
"""⑥calls=6000 达上限 → sleep 桩被调一次、_now 推进到下一窗后 calls 归 0 续跑成功。"""
|
||
clock = {"t": datetime(2026, 7, 16, 12, 0)} # 起始窗 T10（12:00）
|
||
sleeps = []
|
||
|
||
def fake_sleep(secs):
|
||
sleeps.append(secs)
|
||
clock["t"] = clock["t"] + timedelta(seconds=secs) # 睡到下一窗边界（15:00）
|
||
reads = {"2026-07-16T10": (0.0, llm.WINDOW_CALL_CAP), "2026-07-16T15": (0.0, 0)}
|
||
calls = []
|
||
|
||
def chat(prompt, model=None, **kw):
|
||
calls.append(model)
|
||
return "ok", dict(_USAGE)
|
||
bumps = _install(lambda: clock["t"], reads, chat, sleep=fake_sleep)
|
||
content, usage, used = llm.chat_governed("p")
|
||
assert len(sleeps) == 1, f"达上限应只睡一次，实得{len(sleeps)}次"
|
||
assert used == "MiniMax-M3", f"醒来在新窗应 M3 成功，实得{used}"
|
||
assert bumps[0][0] == "2026-07-16T15", f"应记账到下一窗 T15，实得{bumps[0][0]}"
|
||
|
||
|
||
def main():
|
||
tests = [test_window_key, test_seconds_to_next_window, test_cost_usd,
|
||
test_current_quota_policy,
|
||
test_route_first_success, test_route_budget_exhausted,
|
||
test_route_sensitive_then_next, test_route_runtime_then_next,
|
||
test_route_all_fail, test_route_call_cap_sleep]
|
||
for t in tests:
|
||
t()
|
||
print(f" ✓ {t.__name__}")
|
||
print("全部通过")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|