Update in memory client to use FastMCP

src/fastmcp/client/memory.py

@@ -1,11 +1,11 @@
 import contextlib
-from typing import Any, TypeVar
+from typing import TypeVar
 
-from mcp.server import Server
 from mcp.shared.memory import create_connected_server_and_client_session
 from typing_extensions import Unpack
 
 from fastmcp.client.base import BaseClient, ClientKwargs
+from fastmcp.server.server import FastMCP
 
 T = TypeVar("T")
 
@@ -19,33 +19,29 @@ class InMemoryClient(BaseClient):
 
     def __init__(
         self,
-        server: Server[Any],
-        raise_exceptions: bool = False,
+        server: FastMCP,
         **kwargs: Unpack[ClientKwargs],
     ):
         """Initialize an InMemoryClient that connects to an in-memory MCP server.
 
         Args:
-            server: The MCP server instance to connect to
-            raise_exceptions: Whether to raise exceptions from the server
+            server: The FastMCP instance to connect to
             **kwargs: Additional arguments for BaseClient
         """
         super().__init__(**kwargs)
         self.server = server
-        self.raise_exceptions = raise_exceptions
         self._cm_session = None
 
     @contextlib.asynccontextmanager
     async def _connect(self):
         """Set up in-memory connection and session"""
         self._cm_session = create_connected_server_and_client_session(
-            server=self.server,
+            server=self.server._mcp_server,
             read_timeout_seconds=self._read_timeout_seconds,
             sampling_callback=self._sampling_callback,
             list_roots_callback=self._list_roots_callback,
             logging_callback=self._logging_callback,
             message_handler=self._message_handler,
-            raise_exceptions=self.raise_exceptions,
         )
 
         async with self._cm_session as session:

tests/client/__init__.py

@@ -0,0 +1 @@
+"""Client tests package."""

tests/client/test_memory_client.py

@@ -0,0 +1,174 @@
+from typing import cast
+
+import pytest
+from pydantic import AnyUrl
+
+from fastmcp.client.memory import InMemoryClient
+from fastmcp.server.server import FastMCP
+
+
+class _TestException(Exception):
+    """Test exception for testing raise_exceptions behavior."""
+
+    pass
+
+
+@pytest.fixture
+def fastmcp_server():
+    """Fixture that creates a FastMCP server with tools, resources, and prompts."""
+    server = FastMCP("TestServer")
+
+    # Add a tool
+    @server.tool()
+    def greet(name: str) -> str:
+        """Greet someone by name."""
+        return f"Hello, {name}!"
+
+    # Add a tool that raises an exception
+    @server.tool()
+    def error_tool() -> str:
+        """A tool that always raises an exception."""
+        raise _TestException("Deliberate test exception")
+
+    # Add a resource
+    @server.resource(uri="data://users")
+    async def get_users():
+        return ["Alice", "Bob", "Charlie"]
+
+    # Add a resource template
+    @server.resource(uri="data://user/{user_id}")
+    async def get_user(user_id: str):
+        return {"id": user_id, "name": f"User {user_id}", "active": True}
+
+    # Add a prompt
+    @server.prompt()
+    def welcome(name: str) -> str:
+        return f"Welcome to FastMCP, {name}!"
+
+    return server
+
+
+@pytest.mark.asyncio
+async def test_list_tools(fastmcp_server):
+    """Test listing tools with InMemoryClient."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    async with client:
+        result = await client.list_tools()
+
+        # Check that our tools are available
+        assert len(result.tools) == 2
+        tool_names = [tool.name for tool in result.tools]
+        assert "greet" in tool_names
+        assert "error_tool" in tool_names
+
+
+@pytest.mark.asyncio
+async def test_call_tool(fastmcp_server):
+    """Test calling a tool with InMemoryClient."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    async with client:
+        result = await client.call_tool("greet", {"name": "World"})
+
+        # The result content should contain our greeting
+        content_str = str(result.content[0])
+        assert "Hello, World!" in content_str
+
+
+@pytest.mark.asyncio
+async def test_list_resources(fastmcp_server):
+    """Test listing resources with InMemoryClient."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    async with client:
+        result = await client.list_resources()
+
+        # Check that our resource is available
+        assert len(result.resources) == 1
+        assert str(result.resources[0].uri) == "data://users"
+
+
+@pytest.mark.asyncio
+async def test_list_prompts(fastmcp_server):
+    """Test listing prompts with InMemoryClient."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    async with client:
+        result = await client.list_prompts()
+
+        # Check that our prompt is available
+        assert len(result.prompts) == 1
+        assert result.prompts[0].name == "welcome"
+
+
+@pytest.mark.asyncio
+async def test_get_prompt(fastmcp_server):
+    """Test getting a prompt with InMemoryClient."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    async with client:
+        result = await client.get_prompt("welcome", {"name": "Developer"})
+
+        # The result should contain our welcome message
+        result_str = str(result)
+        assert "Welcome to FastMCP, Developer!" in result_str
+
+
+@pytest.mark.asyncio
+async def test_read_resource(fastmcp_server):
+    """Test reading a resource with InMemoryClient."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    async with client:
+        # Use the URI from the resource we know exists in our server
+        uri = cast(
+            AnyUrl, "data://users"
+        )  # Use cast for type hint only, the URI is valid
+        result = await client.read_resource(uri)
+
+        # The contents should include our user list
+        contents_str = str(result.contents[0])
+        assert "Alice" in contents_str
+        assert "Bob" in contents_str
+        assert "Charlie" in contents_str
+
+
+@pytest.mark.asyncio
+async def test_client_connection(fastmcp_server):
+    """Test that the client connects and disconnects properly."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    # Before connection
+    assert not client.is_connected()
+
+    # During connection
+    async with client:
+        assert client.is_connected()
+
+    # After connection
+    assert not client.is_connected()
+
+
+@pytest.mark.asyncio
+async def test_resource_template(fastmcp_server):
+    """Test using a resource template with InMemoryClient."""
+    client = InMemoryClient(server=fastmcp_server)
+
+    async with client:
+        # First, list templates
+        result = await client.list_resource_templates()
+
+        # Check that our template is available
+        assert len(result.resourceTemplates) == 1
+        assert "data://user/{user_id}" in result.resourceTemplates[0].uriTemplate
+
+        # Now use the template with a specific user_id
+        uri = cast(AnyUrl, "data://user/123")
+        result = await client.read_resource(uri)
+
+        # Check the content matches what we expect for the provided user_id
+        content_str = str(result.contents[0])
+        assert '"id": "123"' in content_str
+        assert '"name": "User 123"' in content_str
+        assert '"active": true' in content_str