364 lines
15 KiB
Python
364 lines
15 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 计价(含缓存折扣) /
|
||
模型 max_tokens 主动上限 / Plan 429 不退避与按窗熔断 / 普通 429 重试 /
|
||
chat_governed 六条既有路由（首选成功、预算降级、敏感换模型、不可用换模型、全链失败、调用上限睡窗）。"""
|
||
import os
|
||
import types
|
||
from datetime import datetime, timedelta
|
||
|
||
import muse_llm as llm
|
||
|
||
llm.TOKEN = "test-token"
|
||
|
||
# 直接测 chat() 退避/熔断时允许未治理直连；chat_governed 路径会替换 llm.chat。
|
||
os.environ["MUSE_LLM_ALLOW_UNGOVERNED"] = "1"
|
||
|
||
# 费率缓存预置为兜底表：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
|
||
llm._PLAN_QUOTA_OPEN.clear()
|
||
# 只替换 time.sleep（chat 被打桩后 chat_governed 路径不再用 time.time），不污染真 time 模块
|
||
llm.time = types.SimpleNamespace(sleep=(sleep if sleep else (lambda s: None)), time=lambda: 0)
|
||
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_budget_token_caps():
|
||
"""两条模型策略都在外发前由单次预算推导输出 token 上限。"""
|
||
|
||
previous = llm._PRICING_CACHE
|
||
llm._PRICING_CACHE = dict(llm.PRICING_FALLBACK)
|
||
try:
|
||
content_cap = llm.governed_max_tokens(1.0, 1000)
|
||
opus_cap = llm.fixed_opus_max_tokens(1.0, 1000)
|
||
finally:
|
||
llm._PRICING_CACHE = previous
|
||
assert 0 < content_cap <= 512000
|
||
assert 0 < opus_cap <= llm.FIXED_OPUS_MAX_OUTPUT_TOKENS
|
||
|
||
|
||
def test_current_quota_policy():
|
||
"""额度策略固定为每 5 小时 $24/6000，防止实现和测试再次漂移。"""
|
||
assert llm.WINDOW_BUDGET_USD == 24.0, "当前 MiniMax 每窗预算必须是 $24"
|
||
assert llm.WINDOW_CALL_CAP == 6000, "当前每窗总调用上限必须是 6000"
|
||
|
||
|
||
class _FakeResponse:
|
||
"""只实现 chat() 所需的最小 requests.Response 接口。"""
|
||
|
||
def __init__(self, status_code, text="", data=None):
|
||
self.status_code = status_code
|
||
self.text = text
|
||
self._data = data
|
||
|
||
def raise_for_status(self):
|
||
if self.status_code >= 400:
|
||
raise llm.requests.HTTPError(f"HTTP {self.status_code}: {self.text}")
|
||
|
||
def json(self):
|
||
return self._data
|
||
|
||
|
||
def _success_response():
|
||
return _FakeResponse(200, data={
|
||
"choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}],
|
||
"usage": dict(_USAGE),
|
||
})
|
||
|
||
|
||
def _run_chat_with_responses(model, responses, max_tokens=512000, retries=2):
|
||
"""替换 HTTP 会话并执行一次 chat，返回（结果、请求 payload、sleep 记录）。"""
|
||
posts = []
|
||
sleeps = []
|
||
queue = list(responses)
|
||
|
||
class FakeSession:
|
||
trust_env = True
|
||
|
||
def post(self, url, headers=None, json=None, timeout=None):
|
||
posts.append(dict(json))
|
||
return queue.pop(0)
|
||
|
||
old_session, old_time = llm.requests.Session, llm.time
|
||
try:
|
||
llm.requests.Session = FakeSession
|
||
llm.time = types.SimpleNamespace(time=lambda: 0, sleep=lambda secs: sleeps.append(secs))
|
||
result = llm.chat("p", model=model, max_tokens=max_tokens, retries=retries)
|
||
return result, posts, sleeps
|
||
finally:
|
||
llm.requests.Session, llm.time = old_session, old_time
|
||
|
||
|
||
def test_plan_429_does_not_retry():
|
||
"""Token Plan 429 必须首响应立即抛出，不进入 8/16 秒退避。"""
|
||
body = '{"error":{"message":"已达到 Token Plan 用量上限：请升级套餐"}}'
|
||
posts = []
|
||
sleeps = []
|
||
|
||
class FakeSession:
|
||
trust_env = True
|
||
|
||
def post(self, url, headers=None, json=None, timeout=None):
|
||
posts.append(dict(json))
|
||
return _FakeResponse(429, text=body)
|
||
|
||
old_session, old_time = llm.requests.Session, llm.time
|
||
try:
|
||
llm.requests.Session = FakeSession
|
||
llm.time = types.SimpleNamespace(time=lambda: 0, sleep=lambda secs: sleeps.append(secs))
|
||
try:
|
||
llm.chat("p", model="MiniMax-M3", retries=2)
|
||
raise AssertionError("Plan 429 应抛 PlanQuotaExhausted")
|
||
except llm.PlanQuotaExhausted:
|
||
pass
|
||
finally:
|
||
llm.requests.Session, llm.time = old_session, old_time
|
||
assert len(posts) == 1, f"Plan 429 不应重试，实发 {len(posts)} 次"
|
||
assert sleeps == [], f"Plan 429 不应退避，实睡 {sleeps}"
|
||
|
||
|
||
def test_regular_429_still_retries():
|
||
"""普通速率限制 429 仍按原策略退避重试，不能误触发 Plan 熔断。"""
|
||
result, posts, sleeps = _run_chat_with_responses(
|
||
"MiniMax-M3",
|
||
[_FakeResponse(429, text='{"error":{"message":"rate limit"}}'), _success_response()],
|
||
retries=2,
|
||
)
|
||
assert result[0] == "ok", "普通 429 后应重试成功"
|
||
assert len(posts) == 2, f"普通 429 应重试一次，实发 {len(posts)} 次"
|
||
assert sleeps == [8], f"第一次普通 429 应退避 8 秒，实得 {sleeps}"
|
||
|
||
|
||
def test_model_max_tokens_caps():
|
||
"""只主动裁 M2.7/GLM；M3/deepseek 保留调用方请求值。"""
|
||
cases = [
|
||
("MiniMax-M2.7", 196608),
|
||
("glm-5.2", 12000),
|
||
("MiniMax-M3", 512000),
|
||
("deepseek-v4-flash", 512000),
|
||
]
|
||
for model, expected in cases:
|
||
_, posts, _ = _run_chat_with_responses(model, [_success_response()])
|
||
assert posts[0]["max_tokens"] == expected, \
|
||
f"{model} max_tokens 应为 {expected}，实得 {posts[0]['max_tokens']}"
|
||
|
||
|
||
def test_plan_quota_breaker_same_window_and_recovery():
|
||
"""M3/M2.7 同窗各探一次后跳过；跨窗清空熔断并从 M3 重探。"""
|
||
clock = {"t": datetime(2026, 7, 16, 12, 0)}
|
||
calls = []
|
||
|
||
def chat(prompt, model=None, **kw):
|
||
calls.append((llm.window_key(clock["t"]), model))
|
||
if llm.window_key(clock["t"]) == "2026-07-16T10" and model in llm.MINIMAX_MODELS:
|
||
raise llm.PlanQuotaExhausted("plan exhausted")
|
||
return "ok", dict(_USAGE)
|
||
|
||
_install(lambda: clock["t"], lambda wk: (0.0, 0), chat)
|
||
assert llm.chat_governed("first")[2] == "glm-5.2"
|
||
assert llm.chat_governed("second")[2] == "glm-5.2"
|
||
old_window_models = [m for wk, m in calls if wk == "2026-07-16T10"]
|
||
assert old_window_models == ["MiniMax-M3", "MiniMax-M2.7", "glm-5.2", "glm-5.2"], \
|
||
f"同窗 MiniMax 应各探一次后跳过，实得 {old_window_models}"
|
||
|
||
clock["t"] = datetime(2026, 7, 16, 15, 0)
|
||
assert llm.chat_governed("third")[2] == "MiniMax-M3"
|
||
assert calls[-1] == ("2026-07-16T15", "MiniMax-M3"), "跨窗后应清空熔断并重探 M3"
|
||
|
||
|
||
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():
|
||
"""⑥显式 deadline 覆盖窗等待时，sleep 到下一窗后续跑成功。"""
|
||
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", timeout=4 * 3600)
|
||
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 test_route_call_cap_respects_deadline():
|
||
"""窗等待超过调用总 deadline 时失败关闭，不能无界 sleep。"""
|
||
|
||
clock = {"t": datetime(2026, 7, 16, 12, 0)}
|
||
sleeps = []
|
||
reads = {"2026-07-16T10": (0.0, llm.WINDOW_CALL_CAP)}
|
||
_install(lambda: clock["t"], reads, lambda *_a, **_k: ("ok", dict(_USAGE)),
|
||
sleep=lambda secs: sleeps.append(secs))
|
||
try:
|
||
llm.chat_governed("p", timeout=1)
|
||
except RuntimeError as error:
|
||
assert "deadline" in str(error)
|
||
else:
|
||
raise AssertionError("额度窗等待超过 deadline 必须失败关闭")
|
||
assert sleeps == [], "deadline 不足时不得进入 sleep"
|
||
|
||
|
||
def main():
|
||
tests = [test_window_key, test_seconds_to_next_window, test_cost_usd,
|
||
test_budget_token_caps, test_current_quota_policy,
|
||
test_plan_429_does_not_retry, test_regular_429_still_retries,
|
||
test_model_max_tokens_caps,
|
||
test_plan_quota_breaker_same_window_and_recovery,
|
||
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,
|
||
test_route_call_cap_respects_deadline]
|
||
for t in tests:
|
||
t()
|
||
print(f" ✓ {t.__name__}")
|
||
print("全部通过")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|