zhifu-platform/tests/test_money_budget.py

import asyncio
from datetime import timedelta
from uuid import UUID, uuid4

import pytest
from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401
from test_api_documentation import validate_response
from test_token_budget import sample

from zhifu.core.errors import BusinessError
from zhifu.modules.after_sales.models import BusinessEvent
from zhifu.modules.agent_runtime.budget import change_entry, reserve
from zhifu.modules.agent_runtime.models import TokenEntry
from zhifu.modules.billing.models import BillingChange, MoneyAccount, MoneyEntry
from zhifu.modules.identity.access_models import Principal
from zhifu.modules.identity.auth import now
from zhifu.operations.billing_access import change as access

pytestmark = pytest.mark.integration
PREFIX = "/admin/v1/billing"
PROVIDER = "https://api.deepseek.com"
# 此价格仅用于验证整数运算和事务，不表示供应商报价，不产生供应商消费。
PRICE = {
    "kind": "PRICE",
    "expected_version": 0,
    "provider": PROVIDER,
    "model": "deepseek-flash",
    "input_rate": 2_000_000,
    "cached_rate": 500_000,
    "output_rate": 4_000_000,
    "source": "事务验收专用价格，2026-09-20；不作为供应商账单依据",
    "reason": "验证价格审批与整数金额计算",
}


async def setup(context):
    for actor in ("agent", "reviewer"):
        for permission in ("billing:read", "billing:request", "billing:approve"):
            await access(
                context["app"].state.sessions,
                context["tenant"].bytes,
                context["accounts"][actor].bytes,
                permission,
                True,
                "integration",
                "费用事务验收授权",
            )


async def submit(context, payload):
    path = f"{PREFIX}/requests/{uuid4()}"
    response = await context["call"]("PUT", path, "agent", payload)
    assert response.status_code == 200, response.text
    return path


async def approve(context, payload):
    path = await submit(context, payload)
    response = await context["call"](
        "POST",
        path + "/decisions",
        "reviewer",
        {"expected_version": 1, "approve": True, "reason": "核对完成"},
    )
    assert response.status_code == 200, response.text
    return response.json()


async def budget(context, amount, version=0):
    return await approve(
        context,
        {
            "kind": "BUDGET",
            "expected_version": version,
            "daily_limit": amount,
            "reason": "费用预算事务验收",
        },
    )


async def report(context, suffix="account"):
    response = await context["call"]("GET", PREFIX + "/" + suffix, "agent")
    assert response.status_code == 200, response.text
    return response.json()


async def reserve_sample(context, identifier, tokens=5776, provider=PROVIDER):
    await reserve(
        context["app"].state.sessions,
        context["tenant"].bytes,
        identifier,
        "agent",
        uuid4().bytes,
        sample()["model"],
        tokens,
        provider=provider,
    )


async def test_permissions_idempotency_and_documented_responses(scenario):  # noqa: F811
    call = scenario["call"]
    assert (await call("GET", PREFIX + "/account")).status_code == 403
    assert (await scenario["client"].get(PREFIX + "/account")).status_code == 401
    await setup(scenario)
    path = await submit(scenario, PRICE)
    assert (await call("PUT", path, "agent", PRICE)).status_code == 200
    assert (await call("PUT", path, "agent", {**PRICE, "input_rate": 3})).status_code == 409
    decision = {"expected_version": 1, "approve": True, "reason": "核对完成"}
    assert (await call("POST", path + "/decisions", "agent", decision)).status_code == 403
    results = await asyncio.gather(
        *(call("POST", path + "/decisions", "reviewer", decision) for _ in range(4))
    )
    assert all(response.status_code == 200 for response in results)
    async with scenario["app"].state.sessions() as session:
        count = await session.scalar(
            select(func.count())
            .select_from(BusinessEvent)
            .where(
                BusinessEvent.aggregate_id == UUID(path.rsplit("/", 1)[1]).bytes,
                BusinessEvent.event_type == "billing.approved",
            )
        )
        assert count == 1
    # 外租户员工即使拥有费用权限，也不能使用当前租户的申请标识。
    async with scenario["app"].state.sessions.begin() as session:
        foreign = await session.get(Principal, scenario["accounts"]["foreign"].bytes)
        foreign.kind = "STAFF"
        foreign_tenant = foreign.tenant_id
    await access(
        scenario["app"].state.sessions,
        foreign_tenant,
        scenario["accounts"]["foreign"].bytes,
        "billing:request",
        True,
        "integration",
        "验证租户隔离",
    )
    assert (await call("PUT", path, "foreign", PRICE)).status_code == 404
    for suffix in (
        "account",
        "entries",
        "requests",
        "price?provider=https://api.deepseek.com&model=deepseek-flash",
    ):
        response = await call("GET", PREFIX + "/" + suffix, "agent")
        validate_response(
            scenario["app"].openapi(), PREFIX + "/" + suffix.split("?")[0], "get", response, 200
        )


async def test_atomic_admission_price_snapshot_and_settlement(scenario):  # noqa: F811
    await setup(scenario)
    await budget(scenario, 46_208)
    with pytest.raises(BusinessError) as error:
        await reserve_sample(scenario, uuid4().bytes)
    assert error.value.code == "BILLING_PRICE_MISSING"
    await approve(scenario, PRICE)
    identifiers = [uuid4().bytes for _ in range(8)]
    async with asyncio.timeout(30):
        results = await asyncio.gather(
            *(reserve_sample(scenario, identifier) for identifier in identifiers),
            return_exceptions=True,
        )
    assert sum(result is None for result in results) == 2
    assert all(
        result is None
        or (isinstance(result, BusinessError) and result.code == "MONEY_BUDGET_EXCEEDED")
        for result in results
    )
    accepted = [
        identifier
        for identifier, result in zip(identifiers, results, strict=True)
        if result is None
    ]
    await approve(scenario, {**PRICE, "expected_version": 1, "output_rate": 8_000_000})
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    await asyncio.gather(
        *(change_entry(sessions, tenant, accepted[0], usage=sample()["usage"]) for _ in range(5))
    )
    account = await report(scenario)
    assert account["consumed"] == 912 and account["reserved"] == 23_104
    rows = (await report(scenario, "entries"))["items"]
    assert len(rows) == 2 and all(row["price"]["version"] == 1 for row in rows)
    async with sessions() as session:
        assert (
            await session.scalar(
                select(func.count()).select_from(TokenEntry).where(TokenEntry.tenant_id == tenant)
            )
            == 2
        )


async def test_unknown_reconciliation_and_late_usage(scenario):  # noqa: F811
    await setup(scenario)
    await approve(scenario, PRICE)
    await budget(scenario, 100_000)
    identifier = uuid4().bytes
    await reserve_sample(scenario, identifier)
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    await change_entry(sessions, tenant, identifier)
    assert (await report(scenario))["reserved"] == 23_104
    await change_entry(
        sessions,
        tenant,
        identifier,
        release_reason="供应商未计费记录",
        operator="integration",
        expected_version=2,
    )
    assert (await report(scenario))["reserved"] == 23_104
    entry = (await report(scenario, "entries"))["items"][0]
    await approve(
        scenario,
        {
            "kind": "RECONCILE",
            "entry_id": str(UUID(bytes=identifier)),
            "expected_version": entry["version"],
            "amount": 0,
            "reason": "供应商核对记录：确认未计费",
        },
    )
    assert (await report(scenario))["reserved"] == 0
    await change_entry(sessions, tenant, identifier, usage=sample()["usage"])
    assert (await report(scenario))["consumed"] == 0
    async with sessions() as session:
        discrepancy = await session.scalar(
            select(BusinessEvent).where(
                BusinessEvent.aggregate_id == identifier,
                BusinessEvent.event_type == "billing.usage_discrepancy",
            )
        )
        assert discrepancy.payload == {"calculated": 912, "reconciled": 0}
    # 核实费用优先于迟到回调，再次修正仍需双人审批及账本版本。
    entry = (await report(scenario, "entries"))["items"][0]
    await approve(
        scenario,
        {
            "kind": "RECONCILE",
            "entry_id": str(UUID(bytes=identifier)),
            "expected_version": entry["version"],
            "amount": 912,
            "reason": "供应商最终账单确认记录",
        },
    )
    assert (await report(scenario))["consumed"] == 912


async def test_expiry_revocation_and_concurrent_policy_changes(scenario):  # noqa: F811
    await setup(scenario)
    body = {"kind": "BUDGET", "expected_version": 0, "daily_limit": 100, "reason": "预算核对"}
    paths = [await submit(scenario, body) for _ in range(2)]
    decision = {"expected_version": 1, "approve": True, "reason": "核对完成"}
    results = await asyncio.gather(
        *(scenario["call"]("POST", path + "/decisions", "reviewer", decision) for path in paths)
    )
    assert sorted(response.status_code for response in results) == [200, 409]
    expired = await submit(scenario, {**PRICE})
    async with scenario["app"].state.sessions.begin() as session:
        row = await session.get(BillingChange, UUID(expired.rsplit("/", 1)[1]).bytes)
        row.expires_at = now() - timedelta(seconds=1)
    response = await scenario["call"]("POST", expired + "/decisions", "reviewer", decision)
    assert response.json()["code"] == "BILLING_REQUEST_EXPIRED"
    pending = await submit(scenario, PRICE)
    await access(
        scenario["app"].state.sessions,
        scenario["tenant"].bytes,
        scenario["accounts"]["agent"].bytes,
        "billing:request",
        False,
        "integration",
        "撤销申请权限",
    )
    assert (
        await scenario["call"]("POST", pending + "/decisions", "reviewer", decision)
    ).status_code == 403


async def test_cross_day_account_and_zero_budget(scenario):  # noqa: F811
    await setup(scenario)
    await approve(scenario, PRICE)
    await budget(scenario, 100_000)
    identifier = uuid4().bytes
    await reserve_sample(scenario, identifier)
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    day = now().date() - timedelta(days=1)
    async with sessions.begin() as session:
        session.add(
            MoneyAccount(
                tenant_id=tenant, period=day, daily_limit=100_000, reserved=23_104, consumed=0
            )
        )
        await session.flush()
        row = await session.get(MoneyEntry, identifier)
        row.period = day
        account = await session.get(MoneyAccount, (tenant, now().date()))
        account.reserved = 0
    await budget(scenario, 0, 1)
    await change_entry(sessions, tenant, identifier, usage=sample()["usage"])
    assert (await report(scenario))["consumed"] == 0
    assert (await report(scenario, "account?period=" + day.isoformat()))["consumed"] == 912
    with pytest.raises(BusinessError) as error:
        await reserve_sample(scenario, uuid4().bytes)
    assert error.value.code == "MONEY_BUDGET_EXCEEDED"


async def test_missing_cache_usage_requires_reconciliation(scenario):  # noqa: F811
    await setup(scenario)
    await approve(scenario, PRICE)
    await budget(scenario, 100_000)
    identifier = uuid4().bytes
    await reserve_sample(scenario, identifier)
    # 从已有 usage 移除可选缓存明细，验证证据不足时不会推算费用。
    usage = {key: value for key, value in sample()["usage"].items() if key != "input_token_details"}
    await change_entry(
        scenario["app"].state.sessions, scenario["tenant"].bytes, identifier, usage=usage
    )
    account = await report(scenario)
    assert account["consumed"] == 0 and account["reserved"] == 23_104
    entry = (await report(scenario, "entries"))["items"][0]
    assert entry["status"] == "UNKNOWN" and entry["charged"] is None


async def test_rejection_validation_and_cursor(scenario):  # noqa: F811
    await setup(scenario)
    call = scenario["call"]
    for bad in ({**PRICE, "input_rate": -1}, {**PRICE, "output_rate": 0.5}):
        response = await call("PUT", f"{PREFIX}/requests/{uuid4()}", "agent", bad)
        assert response.status_code == 422
    path = await submit(scenario, PRICE)
    decision = {"expected_version": 1, "approve": False, "reason": "需要重新核实价格"}
    response = await call("POST", path + "/decisions", "reviewer", decision)
    assert response.status_code == 200 and response.json()["status"] == "REJECTED"
    assert (await call("POST", path + "/decisions", "reviewer", decision)).status_code == 200
    response = await call("POST", path + "/decisions", "reviewer", {**decision, "approve": True})
    assert response.status_code == 409
    await submit(scenario, PRICE)
    first = await report(scenario, "requests?limit=1")
    second = await report(scenario, "requests?limit=1&after=" + first["next_cursor"])
    assert first["items"][0]["id"] != second["items"][0]["id"]
    assert second["next_cursor"] is None


async def test_fractional_micro_cny_rounding_and_downward_adjustment(scenario):  # noqa: F811
    await setup(scenario)
    await approve(scenario, {**PRICE, "input_rate": 1, "cached_rate": 1, "output_rate": 1})
    await budget(scenario, 100)
    identifier = uuid4().bytes
    await reserve_sample(scenario, identifier)
    await change_entry(
        scenario["app"].state.sessions,
        scenario["tenant"].bytes,
        identifier,
        usage=sample()["usage"],
    )
    assert (await report(scenario))["consumed"] == 1
    entry = (await report(scenario, "entries"))["items"][0]
    await approve(
        scenario,
        {
            "kind": "RECONCILE",
            "entry_id": str(UUID(bytes=identifier)),
            "expected_version": entry["version"],
            "amount": 0,
            "reason": "供应商确认本次费用减免",
        },
    )
    account = await report(scenario)
    assert account["consumed"] == 0 and account["reserved"] == 0
    assert account["available"] == 100