zhifu-platform/tests/sse_disconnect_acceptance.py

import asyncio
import json
from pathlib import Path
from time import perf_counter
from uuid import uuid4

import httpx
from scale_support import ScaleEnvironment
from sse_capacity_acceptance import wait_released
from test_after_sales import scenario  # noqa: F401


async def test_real_early_disconnect_releases_sse_leases(scenario):  # noqa: F811
    state = scenario["app"].state.state
    assert await state.zcard("sse:global") == 0
    environment = ScaleEnvironment()
    tasks = []
    report = {"maximum_requests": 15, "maximum_concurrency": 12,
              "maximum_seconds": 240, "generation_attempts": 0,
              "maximum_release_wait_seconds": 65,
              "notification_delay_separately_measured": False}
    releases = []
    try:
        async with asyncio.timeout(240):
            await environment.start()
            async with httpx.AsyncClient(
                base_url=environment.base_url, headers=scenario["headers"]["customer"],
                timeout=5, trust_env=False, limits=httpx.Limits(max_connections=12),
            ) as client:
                response = await client.post(
                    "/openapi/v1/conversations", headers={"Idempotency-Key": uuid4().hex},
                    json={"store_id": str(scenario["store"])},
                )
                assert response.status_code == 201
                path = "/openapi/v1/conversations/" + response.json()["conversation_id"] + "/events"
                # 真实 HTTP 返回响应头后立即断开，不等待事件正文。
                async with client.stream("GET", path) as response:
                    assert response.status_code == 200
                started = perf_counter()
                await wait_released(state, 65)
                releases.append(perf_counter() - started)
                await wait_released(state, 65)

                async def interrupted():
                    async with client.stream("GET", path):
                        await asyncio.sleep(5)

                tasks = [asyncio.create_task(interrupted()) for _ in range(12)]
                await asyncio.sleep(0.01)
                for task in tasks:
                    task.cancel()
                await asyncio.gather(*tasks, return_exceptions=True)
                await asyncio.sleep(0.25)
                await wait_released(state, 65)
                async with client.stream("GET", path) as response:
                    assert response.status_code == 200
            await wait_released(state, 65)
            report.update(status="PASSED", remaining_leases=0,
                          maximum_release_seconds=max(releases))
    finally:
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        await environment.close()
        await asyncio.to_thread(
            Path(".work/sse-early-disconnect.json").write_text,
            json.dumps(report, indent=2) + "\n", encoding="utf-8",
        )