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))