mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2026-01-01 21:24:20 +01:00
test_sse_connection is passing
This commit is contained in:
Nick Merrill
@@ -1,173 +0,0 @@
|
||||
import anyio
|
||||
import asyncio
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.routing import Mount, Route
|
||||
import httpx
|
||||
from httpx import ReadTimeout, ASGITransport
|
||||
from starlette.responses import Response
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.server.sse import SseServerTransport
|
||||
from mcp.types import JSONRPCMessage
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def sse_transport():
|
||||
"""Fixture that creates an SSE transport instance."""
|
||||
return SseServerTransport("/messages/")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def sse_app(sse_transport):
|
||||
"""Fixture that creates a Starlette app with SSE endpoints."""
|
||||
async def handle_sse(request):
|
||||
"""Handler for SSE connections."""
|
||||
async def event_generator():
|
||||
try:
|
||||
async with sse_transport.connect_sse(
|
||||
request.scope, request.receive, request._send
|
||||
) as streams:
|
||||
client_to_server, server_to_client = streams
|
||||
# Send initial connection event
|
||||
yield {
|
||||
"event": "endpoint",
|
||||
"data": "/messages",
|
||||
}
|
||||
|
||||
# Process messages
|
||||
async with anyio.create_task_group() as tg:
|
||||
try:
|
||||
async for message in client_to_server:
|
||||
if isinstance(message, Exception):
|
||||
break
|
||||
yield {
|
||||
"event": "message",
|
||||
"data": message.model_dump_json(),
|
||||
}
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
print('cancelled')
|
||||
return
|
||||
except Exception as e:
|
||||
print("unhandled exception:", e)
|
||||
return
|
||||
except Exception:
|
||||
# Log any unexpected errors but allow connection to close gracefully
|
||||
pass
|
||||
|
||||
return EventSourceResponse(event_generator())
|
||||
|
||||
routes = [
|
||||
Route("/sse", endpoint=handle_sse),
|
||||
Mount("/messages", app=sse_transport.handle_post_message),
|
||||
]
|
||||
|
||||
return Starlette(routes=routes)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def test_client(sse_app):
|
||||
"""Create a test client with ASGI transport."""
|
||||
transport = ASGITransport(app=sse_app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport,
|
||||
base_url="http://testserver",
|
||||
timeout=10.0,
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_connection(test_client):
|
||||
"""Test basic SSE connection and message exchange."""
|
||||
async with anyio.create_task_group() as tg:
|
||||
try:
|
||||
async with sse_client(
|
||||
"http://testserver/sse",
|
||||
headers={"Host": "testserver"},
|
||||
timeout=5,
|
||||
sse_read_timeout=5,
|
||||
client=test_client,
|
||||
) as (read_stream, write_stream):
|
||||
# First get the initial endpoint message
|
||||
async with read_stream:
|
||||
init_message = await read_stream.__anext__()
|
||||
assert isinstance(init_message, JSONRPCMessage)
|
||||
|
||||
# Send a test message
|
||||
test_message = JSONRPCMessage.model_validate({"jsonrpc": "2.0", "method": "test"})
|
||||
await write_stream.send(test_message)
|
||||
|
||||
# Receive echoed message
|
||||
async with read_stream:
|
||||
message = await read_stream.__anext__()
|
||||
assert isinstance(message, JSONRPCMessage)
|
||||
assert message.model_dump() == test_message.model_dump()
|
||||
|
||||
# Explicitly close streams
|
||||
await write_stream.aclose()
|
||||
await read_stream.aclose()
|
||||
except Exception as e:
|
||||
pytest.fail(f"Test failed with error: {str(e)}")
|
||||
|
||||
|
||||
# @pytest.mark.anyio
|
||||
# async def test_sse_read_timeout(test_client):
|
||||
# """Test that SSE client properly handles read timeouts."""
|
||||
# with pytest.raises(ReadTimeout):
|
||||
# async with sse_client(
|
||||
# "http://testserver/sse",
|
||||
# headers={"Host": "testserver"},
|
||||
# timeout=5,
|
||||
# sse_read_timeout=2,
|
||||
# client=test_client,
|
||||
# ) as (read_stream, write_stream):
|
||||
# async with read_stream:
|
||||
# # This should timeout since no messages are being sent
|
||||
# await read_stream.__anext__()
|
||||
|
||||
|
||||
# @pytest.mark.anyio
|
||||
# async def test_sse_connection_error(test_client):
|
||||
# """Test SSE client behavior with connection errors."""
|
||||
# with pytest.raises(httpx.HTTPError):
|
||||
# async with sse_client(
|
||||
# "http://testserver/nonexistent",
|
||||
# headers={"Host": "testserver"},
|
||||
# timeout=5,
|
||||
# client=test_client,
|
||||
# ):
|
||||
# pass # Should not reach here
|
||||
|
||||
|
||||
# @pytest.mark.anyio
|
||||
# async def test_sse_multiple_messages(test_client):
|
||||
# """Test sending and receiving multiple SSE messages."""
|
||||
# async with sse_client(
|
||||
# "http://testserver/sse",
|
||||
# headers={"Host": "testserver"},
|
||||
# timeout=5,
|
||||
# sse_read_timeout=5,
|
||||
# client=test_client,
|
||||
# ) as (read_stream, write_stream):
|
||||
# # Send multiple test messages
|
||||
# messages = [
|
||||
# JSONRPCMessage.model_validate({"jsonrpc": "2.0", "method": f"test{i}"})
|
||||
# for i in range(3)
|
||||
# ]
|
||||
|
||||
# for msg in messages:
|
||||
# await write_stream.send(msg)
|
||||
|
||||
# # Receive all echoed messages
|
||||
# received = []
|
||||
# async with read_stream:
|
||||
# for _ in range(len(messages)):
|
||||
# message = await read_stream.__anext__()
|
||||
# assert isinstance(message, JSONRPCMessage)
|
||||
# received.append(message)
|
||||
|
||||
# # Verify all messages were received in order
|
||||
# for sent, received in zip(messages, received):
|
||||
# assert sent.model_dump() == received.model_dump()
|
||||
Reference in New Issue
Block a user