Add basic test for mcp-proxy

pyproject.toml

@@ -11,3 +11,11 @@ build-backend = "setuptools.build_meta"
 
 [project.scripts]
 mcp-proxy = "mcp_proxy.__main__:main"
+
+[tool.pytest.ini_options]
+pythonpath = "src"
+addopts = [
+    "--import-mode=importlib",
+]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"

src/mcp_proxy/__init__.py

@@ -7,7 +7,7 @@ from mcp.client.session import ClientSession
 logger = logging.getLogger(__name__)
 
 
-async def confugure_app(name: str, remote_app: ClientSession):
+async def create_server(name: str, remote_app: ClientSession):
     app = server.Server(name)
 
     async def _list_prompts(_: t.Any) -> types.ServerResult:
@@ -96,6 +96,11 @@ async def confugure_app(name: str, remote_app: ClientSession):
 
     app.request_handlers[types.CompleteRequest] = _complete
 
+    return app
+
+
+async def configure_app(name: str, remote_app: ClientSession):
+    app = await create_server(name, remote_app)
     async with server.stdio_server() as (read_stream, write_stream):
         await app.run(
             read_stream,
@@ -111,4 +116,4 @@ async def run_sse_client(url: str):
         async with ClientSession(read_stream, write_stream) as session:
             response = await session.initialize()
 
-            await confugure_app(response.serverInfo.name, session)
+            await configure_app(response.serverInfo.name, session)

tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for mcp-proxy."""

tests/test_init.py

@@ -0,0 +1,81 @@
+"""Tests for the mcp-proxy module."""
+
+import pytest
+import anyio
+from mcp import types
+from mcp.client.session import ClientSession
+from mcp.server import Server
+from mcp.shared.exceptions import McpError
+from mcp.shared.memory import create_connected_server_and_client_session
+from mcp.server.stdio import stdio_server   
+
+from mcp_proxy import configure_app, create_server
+
+async def run_server() -> None:
+    """Run a stdio server."""
+
+    server = Server("test")
+
+    @server.list_prompts()
+    async def handle_list_prompts() -> list[types.Prompt]:
+        return [
+            types.Prompt(name="prompt1"),
+        ]
+    
+    async with stdio_server() as (read_stream, write_stream):
+        await server.run(read_stream, write_stream, server.create_initialization_options())
+
+
+async def test_list_prompts():
+    """Test list_prompts."""
+
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
+        types.JSONRPCMessage
+    ](1)
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
+        types.JSONRPCMessage
+    ](1)
+
+    server = Server("test")
+
+    @server.list_prompts()
+    async def list_prompts() -> list[types.Prompt]:
+        return [
+            types.Prompt(name="prompt1"),
+        ]
+    
+    async def run_server() -> None:
+        print("running server")
+        await server.run(client_to_server_receive, server_to_client_send, server.create_initialization_options())
+
+    async def listen_session():
+        print("listening session")
+        print(session)
+        async for message in session.incoming_messages:
+            if isinstance(message, Exception):
+                raise message
+            print("message")
+            print(message)
+
+    # Create an in memory connection to the fake server
+    async with create_connected_server_and_client_session(server) as session:
+
+        # Baseline behavior for client
+        result = await session.list_prompts()
+        assert result.prompts == [types.Prompt(name="prompt1")]
+
+        with pytest.raises(McpError, match="Method not found"):
+            await session.list_tools()
+
+        # Create a proxy instance to the in memory server
+        wrapped_server = await create_server("name", session)
+
+        # Create a client to the proxy server
+        async with create_connected_server_and_client_session(server) as wrapped_session:
+            await wrapped_session.initialize()
+
+            result = await wrapped_session.list_prompts()
+            assert result.prompts == [types.Prompt(name="prompt1")]
+
+            with pytest.raises(McpError, match="Method not found"):
+                await wrapped_session.list_tools()