zhifu-platform/tests/test_capacity_quotas.py

import asyncio
from uuid import uuid4

import pytest
from pydantic import ValidationError
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.after_sales.models import BusinessEvent
from zhifu.modules.agent_runtime.capacity import acquire, keys, release
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.worker import claim
from zhifu.modules.memory.worker import claim as claim_memory
from zhifu.operations.tenant_quota import QuotaChange, configure

pytestmark = pytest.mark.integration


@pytest.mark.parametrize("first,maximum", [("interactive", 3), ("background", 2)])
async def test_reserved_capacity_survives_other_class_saturation(scenario, first, maximum):  # noqa: F811
    state = scenario["app"].state.state
    settings = AgentSettings()
    prefix = "verify:capacity:" + uuid4().hex
    other = "background" if first == "interactive" else "interactive"
    identifiers = [uuid4().hex for _ in range(16)]
    try:
        results = await asyncio.gather(
            *(
                acquire(state, settings, first, identifier, prefix=prefix)
                for identifier in identifiers
            )
        )
        assert sum(results) == maximum
        others = await asyncio.gather(
            *(acquire(state, settings, other, uuid4().hex, prefix=prefix) for _ in range(8))
        )
        assert sum(others) == settings.global_concurrency - maximum
        assert await state.zcard(prefix) == settings.global_concurrency
        selected = identifiers[results.index(True)]
        await release(state, selected, prefix=prefix)
        await release(state, selected, prefix=prefix)
        assert await state.zcard(prefix) == settings.global_concurrency - 1
        assert await acquire(state, settings, first, uuid4().hex, prefix=prefix)
        # 使用 Redis 的真实过期分数模拟失联租约到期，下一次领取必须清理全部分类集合。
        for key in keys(prefix)[:3]:
            members = await state.zrange(key, 0, -1)
            await state.zadd(key, dict.fromkeys(members, 0))
        assert await acquire(state, settings, other, uuid4().hex, prefix=prefix)
        assert await state.zcard(prefix) == 1
    finally:
        await state.delete(*keys(prefix))


async def test_quota_update_conflict_audit_and_running_tasks(scenario):  # noqa: F811
    sessions = scenario["app"].state.sessions
    tenant_id = scenario["tenant"].bytes
    before = await configure(sessions, tenant_id)
    assert before["version"] == 0
    change = QuotaChange(
        expected_version=0,
        interactive_limit=1,
        background_limit=0,
        operator="integration",
        reason="验证租户额度调整",
    )
    results = await asyncio.gather(
        configure(sessions, tenant_id, change),
        configure(sessions, tenant_id, change),
        return_exceptions=True,
    )
    assert sum(isinstance(result, ValueError) for result in results) == 1
    assert (await configure(sessions, tenant_id))["version"] == 1
    async with sessions() as session:
        events = (
            await session.scalars(
                select(BusinessEvent).where(
                    BusinessEvent.tenant_id == tenant_id,
                    BusinessEvent.event_type == "tenant.quota_configured",
                )
            )
        ).all()
        assert len(events) == 1
        assert events[0].payload["before"]["version"] == 0
        assert events[0].payload["after"]["interactive_limit"] == 1
    for _ in range(3):
        response = await scenario["call"](
            "POST",
            "/openapi/v1/agent/runs",
            payload={
                "store_id": str(scenario["store"]),
                "message": "查询订单",
            },
        )
        assert response.status_code == 202
    assert await claim(sessions, 2)
    assert await claim(sessions, 2) is None
    await configure(
        sessions,
        tenant_id,
        change.model_copy(update={"expected_version": 1, "interactive_limit": 2}),
    )
    assert await claim(sessions, 2)
    await configure(sessions, tenant_id, change.model_copy(update={"expected_version": 2}))
    assert await claim(sessions, 2) is None
    with pytest.raises(ValueError, match="不存在"):
        await configure(sessions, uuid4().bytes, change)


def test_invalid_reservations_and_quota_are_rejected():
    with pytest.raises(ValidationError):
        AgentSettings(global_concurrency=2, interactive_reserved=2, background_reserved=1)
    with pytest.raises(ValidationError):
        QuotaChange(
            expected_version=0,
            interactive_limit=0,
            background_limit=1,
            operator="operator",
            reason="原因",
        )


async def test_background_quota_pause_and_resume(scenario):  # noqa: F811
    sessions = scenario["app"].state.sessions
    call = scenario["call"]
    response = await call(
        "PUT",
        "/openapi/v1/me/memory/consent",
        payload={"expected_revision": 0, "enabled": True},
    )
    assert response.status_code == 200, response.text
    response = await call(
        "POST", "/openapi/v1/conversations", payload={"store_id": str(scenario["store"])}
    )
    path = "/openapi/v1/conversations/" + response.json()["conversation_id"]
    response = await call(
        "POST",
        path + "/messages",
        payload={"text": "请一直用中文简短回复", "expected_control_version": 1},
    )
    assert response.status_code == 201, response.text
    run_id = response.json()["payload"]["run_id"]
    response = await call("POST", f"/openapi/v1/agent/runs/{run_id}/cancel")
    assert response.status_code == 200
    change = QuotaChange(
        expected_version=0,
        interactive_limit=2,
        background_limit=0,
        operator="integration",
        reason="验证后台暂停",
    )
    await configure(sessions, scenario["tenant"].bytes, change)
    assert await claim_memory(sessions) is None
    await configure(
        sessions,
        scenario["tenant"].bytes,
        change.model_copy(update={"expected_version": 1, "background_limit": 1}),
    )
    assert await claim_memory(sessions)
    assert await claim_memory(sessions) is None


async def test_idle_tenant_reservation_and_oversubscription(scenario):  # noqa: F811
    sessions, state = scenario["app"].state.sessions, scenario["app"].state.state
    settings = AgentSettings()
    tenant = scenario["tenant"]
    change = QuotaChange(
        expected_version=0,
        interactive_limit=2,
        background_limit=1,
        interactive_reserved=2,
        operator="integration",
        reason="验证商家预留总量",
    )
    await configure(sessions, tenant.bytes, change)
    with pytest.raises(ValueError, match="预留总量"):
        await configure(sessions, uuid4().bytes, change)
    prefix = "verify:reserved:" + uuid4().hex
    reserved = {tenant.hex: 2}
    try:
        assert await acquire(
            state,
            settings,
            "interactive",
            "other",
            prefix=prefix,
            tenant="other",
            reserved=reserved,
        )
        assert not await acquire(
            state,
            settings,
            "interactive",
            "excess",
            prefix=prefix,
            tenant="other",
            reserved=reserved,
        )
        assert await acquire(
            state, settings, "background", "memory", prefix=prefix, reserved=reserved
        )
        for identifier in ("reserved-1", "reserved-2"):
            assert await acquire(
                state,
                settings,
                "interactive",
                identifier,
                prefix=prefix,
                tenant=tenant.hex,
                reserved=reserved,
            )
        assert await state.zcard(prefix) == 4
    finally:
        await state.delete(*keys(prefix))