zhifu-platform/tests/sse_capacity_acceptance.py

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)