Skip to content
62 changes: 49 additions & 13 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,20 +57,29 @@
from collections.abc import Iterator

try:
from anyio import create_memory_object_stream, create_task_group, EndOfStream
from mcp.types import (
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
)
from mcp.shared.message import SessionMessage
from httpx import (
ASGITransport,
Request as HttpxRequest,
Response as HttpxResponse,
AsyncByteStream,
AsyncClient,
)
except ImportError:
ASGITransport = None
HttpxRequest = None
HttpxResponse = None
AsyncByteStream = None
AsyncClient = None


try:
from anyio import create_memory_object_stream, create_task_group, EndOfStream
from mcp.types import (
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
)
from mcp.shared.message import SessionMessage
except ImportError:
create_memory_object_stream = None
create_task_group = None
Expand All @@ -81,12 +90,6 @@
JSONRPCRequest = None
SessionMessage = None

ASGITransport = None
HttpxRequest = None
HttpxResponse = None
AsyncByteStream = None
AsyncClient = None


SENTRY_EVENT_SCHEMA = "./checkouts/data-schemas/relay/event.schema.json"

Expand Down Expand Up @@ -1013,6 +1016,39 @@ async def inner(values):
return inner


@pytest.fixture
def server_side_event_chunks():
def inner(events):
for event in events:
payload = event.model_dump()
chunk = f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n"
yield chunk.encode("utf-8")

return inner


@pytest.fixture
def get_model_response():
def inner(response_content, serialize_pydantic=False):
model_request = HttpxRequest(
"POST",
"/responses",
)

if serialize_pydantic:
response_content = json.dumps(response_content.model_dump()).encode("utf-8")

response = HttpxResponse(
200,
request=model_request,
content=response_content,
)

return response

return inner


class MockServerRequestHandler(BaseHTTPRequestHandler):
def do_GET(self): # noqa: N802
# Process an HTTP GET request and return a response.
Expand Down
Loading
Loading