mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 06:54:18 +01:00
69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
from unittest.mock import patch
|
|
|
|
import anyio
|
|
import pytest
|
|
|
|
from mcp.shared.session import BaseSession
|
|
from mcp.types import (
|
|
ClientRequest,
|
|
EmptyResult,
|
|
PingRequest,
|
|
)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_send_request_stream_cleanup():
|
|
"""
|
|
Test that send_request properly cleans up streams when an exception occurs.
|
|
|
|
This test mocks out most of the session functionality to focus on stream cleanup.
|
|
"""
|
|
|
|
# Create a mock session with the minimal required functionality
|
|
class TestSession(BaseSession):
|
|
async def _send_response(self, request_id, response):
|
|
pass
|
|
|
|
# Create streams
|
|
write_stream_send, write_stream_receive = anyio.create_memory_object_stream(1)
|
|
read_stream_send, read_stream_receive = anyio.create_memory_object_stream(1)
|
|
|
|
# Create the session
|
|
session = TestSession(
|
|
read_stream_receive,
|
|
write_stream_send,
|
|
object, # Request type doesn't matter for this test
|
|
object, # Notification type doesn't matter for this test
|
|
)
|
|
|
|
# Create a test request
|
|
request = ClientRequest(
|
|
PingRequest(
|
|
method="ping",
|
|
)
|
|
)
|
|
|
|
# Patch the _write_stream.send method to raise an exception
|
|
async def mock_send(*args, **kwargs):
|
|
raise RuntimeError("Simulated network error")
|
|
|
|
# Record the response streams before the test
|
|
initial_stream_count = len(session._response_streams)
|
|
|
|
# Run the test with the patched method
|
|
with patch.object(session._write_stream, "send", mock_send):
|
|
with pytest.raises(RuntimeError):
|
|
await session.send_request(request, EmptyResult)
|
|
|
|
# Verify that no response streams were leaked
|
|
assert len(session._response_streams) == initial_stream_count, (
|
|
f"Expected {initial_stream_count} response streams after request, "
|
|
f"but found {len(session._response_streams)}"
|
|
)
|
|
|
|
# Clean up
|
|
await write_stream_send.aclose()
|
|
await write_stream_receive.aclose()
|
|
await read_stream_send.aclose()
|
|
await read_stream_receive.aclose()
|