zhifu-platform/tests/money_budget_model_acceptance.py

from uuid import UUID

import pytest
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401
from test_money_budget import PRICE, approve, budget, setup

from zhifu.modules.agent_runtime.budget_callback import CallAllowance
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.models import AgentRun, ModelAttempt
from zhifu.modules.agent_runtime.worker import execute_one
from zhifu.operations.token_budget import entries


@pytest.mark.parametrize("code", ["MONEY_BUDGET_EXCEEDED", "BILLING_PRICE_MISSING"])
async def test_money_admission_prevents_model_send(scenario, code):  # noqa: F811
    """使用已配置的真实客户端，预算回调在供应商请求前拒绝；不发送生成请求。"""
    await setup(scenario)
    await budget(scenario, 0)
    if code == "MONEY_BUDGET_EXCEEDED":
        await approve(scenario, PRICE)
    sessions = scenario["app"].state.sessions
    state = scenario["app"].state.state
    async for _ in state.scan_iter(match="monitor:worker:agent:*"):
        raise RuntimeError("请停止 Agent 工作进程并等待心跳过期")
    response = await scenario["call"](
        "POST",
        "/openapi/v1/agent/runs",
        payload={"store_id": str(scenario["store"]), "message": "你好"},
    )
    assert response.status_code == 202
    identifier = UUID(response.json()["run_id"]).bytes
    assert await execute_one(
        sessions, state, AgentSettings(max_output_tokens=128), allowance=CallAllowance(1)
    )
    async with sessions() as session:
        run = await session.get(AgentRun, identifier)
        attempt = await session.scalar(
            select(ModelAttempt).where(ModelAttempt.run_id == identifier)
        )
        assert run.status == "FAILED" and run.error_code == code
        assert attempt.status == "NOT_SENT"
    assert (await entries(sessions, scenario["tenant"].bytes))["items"] == []