import asyncio import hashlib import json import os import secrets from datetime import timedelta from pathlib import Path from statistics import median, quantiles from time import perf_counter from uuid import uuid4 import httpx from sqlalchemy import func, select from test_after_sales import scenario # noqa: F401 from zhifu.modules.agent_runtime.models import ModelAttempt from zhifu.modules.identity.access_models import AccessSession, Principal, StoreGrant from zhifu.modules.identity.auth import now async def wait_released(state, maximum_seconds=10): if not 1 <= maximum_seconds <= 65: raise ValueError("SSE 额度恢复等待上限必须在 1–65 秒之间") # 默认验证主动释放;代理断开验证可明确等待既有异常租约期限。 for _ in range(maximum_seconds * 10): if await state.zcard("sse:global") == 0: return await asyncio.sleep(0.1) raise AssertionError("SSE 断开后额度未在指定期限内释放") async def test_two_hundred_real_sse_connections(scenario): # noqa: F811 sessions, state = scenario["app"].state.sessions, scenario["app"].state.state assert await state.zcard("sse:global") == 0, "需要独立的 SSE 验收环境" identities = [] # 预置67个独立验收账户及有效会话;每次HTTP请求仍执行真实服务端认证。 async with sessions.begin() as session: customer = await session.get(Principal, scenario["accounts"]["customer"].bytes) for index in range(67): actor, token = uuid4(), secrets.token_urlsafe(32) session.add( Principal( id=actor.bytes, tenant_id=scenario["tenant"].bytes, username=f"sse-{index}", password_hash=customer.password_hash, kind="CUSTOMER", ) ) await session.flush() session.add( StoreGrant( tenant_id=scenario["tenant"].bytes, principal_id=actor.bytes, store_id=scenario["store"].bytes, permission="orders:read", ) ) session.add( AccessSession( token_hash=hashlib.sha256(token.encode()).digest(), tenant_id=scenario["tenant"].bytes, principal_id=actor.bytes, expires_at=now() + timedelta(minutes=10), ) ) identities.append({"Authorization": "Bearer " + token}) ready = asyncio.Queue() finish = asyncio.Event() latencies = [] statuses = [] tasks = [] paths = [] started = perf_counter() limits = httpx.Limits(max_connections=205, max_keepalive_connections=205) async with httpx.AsyncClient( base_url=os.environ.get("ACCEPTANCE_API_URL", "http://127.0.0.1:18000"), limits=limits, timeout=30, ) as client: async def connect(path, headers): begin = perf_counter() async with client.stream("GET", path + "/events", headers=headers) as response: statuses.append(response.status_code) latencies.append((perf_counter() - begin) * 1000) error = (await response.aread()).decode() if response.status_code != 200 else "" await ready.put((response.status_code, error)) assert response.status_code == 200, error assert response.headers["content-type"].startswith("text/event-stream") await finish.wait() try: # 请求上限271次:67次创建、200次订阅、2次拒绝、1次恢复订阅及1次健康检查。 async with asyncio.timeout(90): for headers in identities: response = await client.post( "/openapi/v1/conversations", headers={**headers, "Idempotency-Key": str(uuid4())}, json={"store_id": str(scenario["store"])}, ) assert response.status_code == 201, response.text paths.append("/openapi/v1/conversations/" + response.json()["conversation_id"]) for index in range(200): tasks.append( asyncio.create_task(connect(paths[index // 3], identities[index // 3])) ) for _ in range(200): status, error = await ready.get() assert status == 200, error assert await state.zcard("sse:global") == 200 # 最后一个账户仅占两个连接,新增请求因全局200连接上限被拒绝。 response = await client.get(paths[-1] + "/events", headers=identities[-1]) assert response.status_code == 429 and response.json()["code"] == "SSE_LIMIT" response = await client.get(paths[0] + "/events", headers=identities[0]) assert response.status_code == 429 await asyncio.sleep(10) assert await state.zcard("sse:global") == 200 assert (await client.get("/health/ready")).status_code == 200 finish.set() await asyncio.gather(*tasks) await wait_released(state) async with client.stream( "GET", paths[0] + "/events", headers=identities[0] ) as response: assert response.status_code == 200 async with sessions() as session: attempts = await session.scalar( select(func.count()) .select_from(ModelAttempt) .where( ModelAttempt.tenant_id == scenario["tenant"].bytes, ) ) assert attempts == 0 percentiles = quantiles(latencies, n=100) report = { "connections": 200, "accounts": 67, "request_limit": 271, "duration_limit_seconds": 90, "held_seconds": 10, "seconds": perf_counter() - started, "http_200": statuses.count(200), "http_429": 2, "p50_ms": median(latencies), "p95_ms": percentiles[94], "p99_ms": percentiles[98], "generation_attempts": attempts, "disconnect_recovery": True, } await asyncio.to_thread( Path(".work/sse-capacity.json").write_text, json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) finally: finish.set() for task in tasks: if not task.done(): task.cancel() await asyncio.gather(*tasks, return_exceptions=True) # 等待断开清理完成,确保fixture删除账户前服务端结束查询。 await wait_released(state)