add tests using fastmcp

Author: brs96

mcp_server/pyproject.toml

@@ -23,4 +23,5 @@ target-version = "py311"
 dev = [
     "ruff>=0.12.2",
     "pytest>=8.4.1",
+    "pytest-asyncio>=1.0.0",
 ]

mcp_server/src/mcp_server_neo4j_gds/server.py

@@ -44,28 +44,62 @@ def serialize_result(result: Any) -> str:
 
 
 def create_algorithm_tool(mcp: FastMCP, tool_name: str, gds: GraphDataScience):
-    """Create and register an algorithm tool with proper closure"""
-
-    async def algorithm_tool(parameters: Dict[str, Any] = None) -> str:
-        """Execute algorithm tool with parameters dictionary"""
-        try:
-            handler = AlgorithmRegistry.get_handler(tool_name, gds)
-            result = handler.execute(parameters or {})
-            return serialize_result(result)
-        except Exception as e:
-            return f"Error executing {tool_name}: {str(e)}"
-
-    # Set the function name and docstring
-    algorithm_tool.__name__ = tool_name
-    algorithm_tool.__doc__ = f"Execute {tool_name} algorithm with parameters dictionary"
+    """Create and register an algorithm tool using a single parameters dict approach"""
+    handler = AlgorithmRegistry.get_handler(tool_name, gds)
+
+    # Get tool definition from specs to extract parameter information
+    all_tool_definitions = (
+        centrality_tool_definitions
+        + community_tool_definitions
+        + path_tool_definitions
+        + similarity_tool_definitions
+    )
 
-    # Register the tool with the server
-    mcp.tool(algorithm_tool)
+    tool_def = None
+    for tool in all_tool_definitions:
+        if tool.name == tool_name:
+            tool_def = tool
+            break
+
+    if tool_def:
+        # Create a function that accepts a single parameters dict
+        # This is the only approach that works consistently with FastMCP
+        async def algorithm_tool(parameters: Dict[str, Any] = None) -> str:
+            """Execute algorithm tool with parameters dictionary"""
+            try:
+                result = handler.execute(parameters or {})
+                return serialize_result(result)
+            except Exception as e:
+                return f"Error executing {tool_name}: {str(e)}"
+
+        # Set the function name and docstring
+        algorithm_tool.__name__ = tool_name
+        algorithm_tool.__doc__ = tool_def.description
+
+        # Register the tool with the server
+        mcp.tool(algorithm_tool)
+    else:
+        # Fallback for tools without specs
+        async def algorithm_tool(parameters: Dict[str, Any] = None) -> str:
+            """Execute algorithm tool with parameters dictionary"""
+            try:
+                result = handler.execute(parameters or {})
+                return serialize_result(result)
+            except Exception as e:
+                return f"Error executing {tool_name}: {str(e)}"
+
+        algorithm_tool.__name__ = tool_name
+        algorithm_tool.__doc__ = (
+            f"Execute {tool_name} algorithm with parameters dictionary"
+        )
+        mcp.tool(algorithm_tool)
 
 
-def main(db_url: str, username: str, password: str, database: str = None):
-    """Main function that sets up and runs the FastMCP server"""
-    logger.info(f"Starting MCP Server for {db_url} with username {username}")
+def setup_server(
+    db_url: str, username: str, password: str, database: str = None
+) -> FastMCP:
+    """Set up the FastMCP server with all tools registered"""
+    logger.info(f"Setting up MCP Server for {db_url} with username {username}")
     if database:
         logger.info(f"Connecting to database: {database}")
 
@@ -142,6 +176,12 @@ async def list_tools() -> str:
     for tool_name in algorithm_names:
         create_algorithm_tool(mcp, tool_name, gds)
 
+    return mcp
+
+
+def main(db_url: str, username: str, password: str, database: str = None):
+    """Main function that sets up and runs the FastMCP server"""
+    mcp = setup_server(db_url, username, password, database)
     # Run the server - FastMCP will handle the event loop
     mcp.run()
 

mcp_server/tests/conftest.py

@@ -0,0 +1,14 @@
+import pytest
+import asyncio
+
+
+# Configure pytest-asyncio
+pytest_plugins = ["pytest_asyncio"]
+
+
+@pytest.fixture(scope="session")
+def event_loop():
+    """Create an instance of the default event loop for the test session."""
+    loop = asyncio.get_event_loop_policy().new_event_loop()
+    yield loop
+    loop.close()

mcp_server/tests/test_basic.py

@@ -1 +1,65 @@
+import pytest
+import asyncio
+import json
+import os
+from dotenv import load_dotenv
+from fastmcp import Client
 
+
+@pytest.fixture
+def mcp_server():
+    """Create a FastMCP server instance using the actual server setup"""
+    # Load environment variables
+    load_dotenv("../../../.env")
+
+    # Import the actual server setup function
+    from mcp_server_neo4j_gds.server import setup_server
+
+    # Get environment variables
+    db_url = os.environ.get("NEO4J_URI")
+    username = os.environ.get("NEO4J_USERNAME", "neo4j")
+    password = os.environ.get("NEO4J_PASSWORD")
+    database = os.environ.get("NEO4J_DATABASE")
+
+    # Use the actual server setup function
+    server = setup_server(db_url, username, password, database)
+
+    return server
+
+
+@pytest.mark.asyncio
+async def test_find_shortest_path(mcp_server):
+    """Test the find_shortest_path tool with correct parameters using the actual server"""
+    # Pass the server directly to the Client constructor for in-memory testing
+    async with Client(mcp_server) as client:
+        # Test the find_shortest_path tool with correct parameter names
+        result = await client.call_tool(
+            "find_shortest_path", {"start_node": "Tower Hill", "end_node": "Paddington"}
+        )
+
+        # Parse the result
+        result_data = json.loads(result.data)
+        breakpoint()
+
+        # Assertions
+        assert "totalCost" in result_data, "Result should contain totalCost"
+        assert "nodeIds" in result_data, "Result should contain nodeIds"
+        assert "nodeNames" in result_data, "Result should contain nodeNames"
+        assert "path" in result_data, "Result should contain path"
+        assert "costs" in result_data, "Result should contain costs"
+
+        # Check that we got a valid path
+        assert result_data["totalCost"] > 0, "Total cost should be positive"
+        assert len(result_data["nodeIds"]) > 0, "Should have at least one node in path"
+        assert len(result_data["nodeNames"]) > 0, "Should have at least one node name"
+
+        print(f"✅ find_shortest_path test passed!")
+        print(
+            f"Path from {result_data['nodeNames'][0]} to {result_data['nodeNames'][-1]}"
+        )
+        print(f"Total cost: {result_data['totalCost']}")
+        print(f"Number of nodes: {len(result_data['nodeIds'])}")
+
+
+if __name__ == "__main__":
+    asyncio.run(test_find_shortest_path())