from contextlib import asynccontextmanager

from unittest.mock import patch

import pytest

import mcp.shared.memory

from mcp.shared.message import SessionMessage

from mcp.types import (

JSONRPCNotification,

JSONRPCRequest,

)

class SpyMemoryObjectSendStream:

def __init__(self, original_stream):

self.original_stream = original_stream

self.sent_messages: list[SessionMessage] = []

async def send(self, message):

self.sent_messages.append(message)

await self.original_stream.send(message)

async def aclose(self):

await self.original_stream.aclose()

async def __aenter__(self):

return self

async def __aexit__(self, *args):

await self.aclose()

class StreamSpyCollection:

def __init__(

self,

client_spy: SpyMemoryObjectSendStream,

server_spy: SpyMemoryObjectSendStream,

):

self.client = client_spy

self.server = server_spy

def clear(self) -> None:

"""Clear all captured messages."""

self.client.sent_messages.clear()

self.server.sent_messages.clear()

def get_client_requests(self, method: str | None = None) -> list[JSONRPCRequest]:

"""Get client-sent requests, optionally filtered by method."""

return [

req.message.root

for req in self.client.sent_messages

if isinstance(req.message.root, JSONRPCRequest) and (method is None or req.message.root.method == method)

]

def get_server_requests(self, method: str | None = None) -> list[JSONRPCRequest]:

"""Get server-sent requests, optionally filtered by method."""

return [

req.message.root

for req in self.server.sent_messages

if isinstance(req.message.root, JSONRPCRequest) and (method is None or req.message.root.method == method)

]

def get_client_notifications(self, method: str | None = None) -> list[JSONRPCNotification]:

"""Get client-sent notifications, optionally filtered by method."""

return [

notif.message.root

for notif in self.client.sent_messages

if isinstance(notif.message.root, JSONRPCNotification)

and (method is None or notif.message.root.method == method)

]

def get_server_notifications(self, method: str | None = None) -> list[JSONRPCNotification]:

"""Get server-sent notifications, optionally filtered by method."""

return [

notif.message.root

for notif in self.server.sent_messages

if isinstance(notif.message.root, JSONRPCNotification)

and (method is None or notif.message.root.method == method)

]

@pytest.fixture

def stream_spy():

"""Fixture that provides spies for both client and server write streams.

Example usage:

async def test_something(stream_spy):

# ... set up server and client ...

spies = stream_spy()

# Run some operation that sends messages

await client.some_operation()

# Check the messages

requests = spies.get_client_requests(method="some/method")

assert len(requests) == 1

# Clear for the next operation

spies.clear()

"""

client_spy = None

server_spy = None

# Store references to our spy objects

def capture_spies(c_spy, s_spy):

nonlocal client_spy, server_spy

client_spy = c_spy

server_spy = s_spy

# Create patched version of stream creation

original_create_streams = mcp.shared.memory.create_client_server_memory_streams

@asynccontextmanager

async def patched_create_streams():

async with original_create_streams() as (client_streams, server_streams):

client_read, client_write = client_streams

server_read, server_write = server_streams

# Create spy wrappers

spy_client_write = SpyMemoryObjectSendStream(client_write)

spy_server_write = SpyMemoryObjectSendStream(server_write)

# Capture references for the test to use

capture_spies(spy_client_write, spy_server_write)

yield (client_read, spy_client_write), (server_read, spy_server_write)

# Apply the patch for the duration of the test

with patch("mcp.shared.memory.create_client_server_memory_streams", patched_create_streams):

# Return a collection with helper methods

def get_spy_collection() -> StreamSpyCollection:

assert client_spy is not None, "client_spy was not initialized"

assert server_spy is not None, "server_spy was not initialized"

return StreamSpyCollection(client_spy, server_spy)

yield get_spy_collection