muse-agent-example/.claude/skills/call-content-model/scripts/test_quota.py

#!/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 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
    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_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():
    """⑥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_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]
    for t in tests:
        t()
        print(f"  ✓ {t.__name__}")
    print("全部通过")


if __name__ == "__main__":
    main()