wehub-resource-sync
85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
"""
|
|
Shared SSE-collection helpers for server integration tests.
|
|
|
|
The integration tests routinely connect to ``/v1/responses`` with
|
|
``stream=True`` and need to assert on the *sequence* of SSE events.
|
|
``CapturedEvent`` formalizes the ``(event_type, parsed_data)`` shape
|
|
the suite has been using ad-hoc, and ``collect_sse_events()`` runs
|
|
the full event stream into a list — including the synthetic
|
|
``"done"`` marker for the trailing ``[DONE]`` line.
|
|
|
|
Required by the concurrency tests. Lives here rather than
|
|
``conftest.py`` because these are plain helpers, not pytest
|
|
fixtures.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from httpx_sse import aconnect_sse
|
|
|
|
|
|
@dataclass
|
|
class CapturedEvent:
|
|
"""
|
|
One SSE event captured from a streaming response.
|
|
|
|
:param event: The SSE ``event:`` field, e.g.
|
|
``"response.created"``, ``"response.output_text.delta"``,
|
|
``"response.completed"``, or the synthetic ``"done"`` marker
|
|
for the trailing ``[DONE]`` line.
|
|
:param data: The parsed ``data:`` payload. A dict for normal
|
|
events, the literal string ``"[DONE]"`` for the terminal
|
|
marker.
|
|
"""
|
|
|
|
event: str
|
|
# dict[str, Any] | str — most events are JSON dicts, the [DONE]
|
|
# marker is a bare string.
|
|
data: dict[str, Any] | str
|
|
|
|
|
|
async def collect_sse_events(
|
|
client: httpx.AsyncClient,
|
|
*,
|
|
method: str,
|
|
url: str,
|
|
json_body: dict[str, Any],
|
|
) -> list[CapturedEvent]:
|
|
"""
|
|
Connect to an SSE endpoint and collect every event into a list.
|
|
|
|
Used by tests that assert on event ordering, content, or counts.
|
|
The trailing ``[DONE]`` line is captured as a ``CapturedEvent``
|
|
with ``event="done"`` and ``data="[DONE]"`` so tests can
|
|
distinguish a clean stream close from a mid-stream disconnect.
|
|
|
|
:param client: The HTTP client to use, typically the test
|
|
``httpx.AsyncClient``.
|
|
:param method: HTTP method, e.g. ``"POST"``.
|
|
:param url: Target URL, e.g. ``"/v1/responses"``.
|
|
:param json_body: JSON body for the request, e.g.
|
|
``{"model": "test-agent", "input": "Hi", "stream": True}``.
|
|
:returns: List of :class:`CapturedEvent` in arrival order. An
|
|
empty list means the server closed the stream without
|
|
emitting any events (likely an error before the SSE
|
|
handshake).
|
|
"""
|
|
events: list[CapturedEvent] = []
|
|
async with aconnect_sse(
|
|
client,
|
|
method,
|
|
url,
|
|
json=json_body,
|
|
) as event_source:
|
|
async for sse in event_source.aiter_sse():
|
|
if sse.data == "[DONE]":
|
|
events.append(CapturedEvent(event="done", data="[DONE]"))
|
|
else:
|
|
events.append(CapturedEvent(event=sse.event, data=json.loads(sse.data)))
|
|
return events
|