games-development-ai/cheap-worker/tests/test_budget_gate.py

"""
test_budget_gate.py — 便宜档预算闸 <¥10 fail-closed 接线单测(M1 U4)。

直接测 CircuitBreakerMiddleware 在便宜档参数(enable_rmb_gate=True / rmb_hard_limit=10.0)下的 on_model_call 行为,
全程 mock(传 pricing_params 让 ¥ 闸激活但不触网、不取价、不真跑模型),验证:
  · ¥10 越线 fail-closed 抛 Tier2CircuitBreak(budget);正常生成不被误拦。
  · 取价不可达降级(_rmb_gate_active=False)时放行(由 on_reply 次数闸兜),不静默阻断。
  · enable_rmb_gate=False 总开关旁路。
  · 便宜档 ¥10 与 tier2 默认口径互不串。
成本落盘字段(costRmb/rmbGate/tokens)的正确性由 U4 的一次真实生成验证(plan U4 Verification),非本单测覆盖面。

跑:cheap-worker/.venv/bin/python cheap-worker/tests/test_budget_gate.py
"""

import asyncio
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # → cheap-worker/
import _bootstrap  # noqa: E402,F401  仅加 sys.path(import 时不取 key、不触网)
from worker.middleware import CircuitBreakerMiddleware, Tier2CircuitBreak  # noqa: E402

# 非 None 的计费参数 → reset() 里 _rmb_gate_active=True、且 _pricing_fetch_tried=True(on_model_call 不再去取价、不触网)。
# 具体费率不重要:测试让 _estimate_call_rmb 走「均价路径」(预设 spent_rmb>0 + model_calls>0),不进 cost.compute。
_FAKE_PRICING = {"pricing": {}, "qpu": 1, "usd_rate": 1.0}


class _FakeModel:
    """on_model_call 从 input_kwargs['current_model'].model 取模型名。"""
    model = "MiniMax-M3"


async def _fake_handler(**_kwargs):
    """mock 裸模型调用:返回一个非 async-gen、usage=None 的响应(不触发折算累加)。"""
    class _Resp:
        usage = None
    return _Resp()


def _cheap_breaker(**kw):
    """便宜档参数 breaker(¥ 闸开、¥10、带 fake 计费参数 → 金额闸激活不触网)。"""
    return CircuitBreakerMiddleware(enable_rmb_gate=True, rmb_hard_limit=10.0,
                                    pricing_params=_FAKE_PRICING, **kw)


def _call(breaker):
    """同步触发一次 on_model_call(内部 asyncio.run);抛 Tier2CircuitBreak 则透传。"""
    return asyncio.run(breaker.on_model_call(None, {"current_model": _FakeModel()}, _fake_handler))


# ───────────────────────── ¥ 闸 fail-closed / 不误拦 ─────────────────────────

def test_fail_closed_over_10():
    """已花 ¥9.99 + 本次预估(均价 9.99/2≈5)> ¥10 → fail-closed 抛 budget。"""
    b = _cheap_breaker()
    b.spent_rmb = 9.99
    b.model_calls = 2  # 均价 = spent_rmb/model_calls ≈ 5.0 → 9.99+5 > 10
    raised = None
    try:
        _call(b)
    except Tier2CircuitBreak as e:
        raised = e
    assert raised is not None, "越 ¥10 应 fail-closed 抛 Tier2CircuitBreak"
    assert raised.kind == "budget", f"熔断 kind 应为 budget,实际 {raised.kind}"


def test_normal_not_tripped():
    """已花 ¥1 + 本次预估(均价 0.5)< ¥10 → 正常放行、不误拦。"""
    b = _cheap_breaker()
    b.spent_rmb = 1.0
    b.model_calls = 2  # 均价 0.5 → 1.0+0.5 = 1.5 < 10
    res = _call(b)
    assert res is not None, "正常额度内应放行返回响应"


def test_boundary_just_under_10():
    """边界:已花 ¥7 + 预估(均价 7/3≈2.33)≈ 9.33 < 10 → 不拦。"""
    b = _cheap_breaker()
    b.spent_rmb = 7.0
    b.model_calls = 3  # 均价 ≈ 2.33 → 9.33 < 10
    res = _call(b)
    assert res is not None


# ───────────────────────── 取价降级 / 总开关 ─────────────────────────

def test_degraded_passes_through_no_silent_block():
    """取价不可达降级(_rmb_gate_active=False)→ ¥ 闸不硬拦、放行(由次数闸兜),不静默阻断。"""
    b = _cheap_breaker()
    b._rmb_gate_active = False  # 模拟取价失败后的降级态
    b.spent_rmb = 100.0         # 即便"看似超支",降级时金额闸不生效、不硬拦(次数闸在 on_reply 兜)
    res = _call(b)
    assert res is not None, "降级态应放行(不静默阻断、由次数闸兜底)"


def test_disabled_gate_bypasses():
    """enable_rmb_gate=False 总开关 → on_model_call 直接放行(回归保护开关语义)。"""
    b = CircuitBreakerMiddleware(enable_rmb_gate=False, rmb_hard_limit=10.0)
    b.spent_rmb = 100.0
    res = _call(b)
    assert res is not None, "关闸时应直接放行"


# ───────────────────────── 便宜档 ¥10 与 tier2 口径隔离 ─────────────────────────

def test_cheap_limit_10_isolated_from_tier2_default():
    """便宜档实例 ¥10 显式传入;tier2 默认实例走 genconfig 默认(≠10),两者各自独立、互不串。"""
    cheap = CircuitBreakerMiddleware(enable_rmb_gate=True, rmb_hard_limit=10.0)
    assert cheap.rmb_hard_limit == 10.0 and cheap.enable_rmb_gate is True
    tier2_default = CircuitBreakerMiddleware()  # 不传 rmb_hard_limit → genconfig 默认(tier2 ¥3 口径)
    assert tier2_default.rmb_hard_limit != 10.0, "便宜档 ¥10 不应污染 tier2 默认上限"
    assert cheap.rmb_hard_limit == 10.0, "便宜档实例上限仍为 ¥10"


if __name__ == "__main__":
    _fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)]
    _failed = 0
    for _fn in _fns:
        try:
            _fn()
            print(f"  PASS {_fn.__name__}")
        except Exception as e:  # noqa: BLE001
            _failed += 1
            print(f"  FAIL {_fn.__name__}: {type(e).__name__}: {e}")
    print(f"\n{len(_fns) - _failed}/{len(_fns)} passed")
    sys.exit(1 if _failed else 0)