From 5bc07b0c5dbd315a31b5701a629cd6301120d330 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 20 Mar 2025 22:03:56 -0700
Subject: [PATCH] test tool registry

---
 .../mcp_server/test_tool_registry.py          | 95 +++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 tests/litellm/proxy/experimental/mcp_server/test_tool_registry.py

diff --git a/tests/litellm/proxy/experimental/mcp_server/test_tool_registry.py b/tests/litellm/proxy/experimental/mcp_server/test_tool_registry.py
new file mode 100644
index 00000000000..d5ba9744c7d
--- /dev/null
+++ b/tests/litellm/proxy/experimental/mcp_server/test_tool_registry.py
@@ -0,0 +1,95 @@
+import json
+import os
+import sys
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.proxy._experimental.mcp_server.tool_registry import MCPToolRegistry
+
+
+# Test handler function
+def example_handler(input_data):
+    return {"result": input_data}
+
+
+def test_register_and_get_tool():
+    registry = MCPToolRegistry()
+
+    # Test registering a tool
+    registry.register_tool(
+        name="test_tool",
+        description="A test tool",
+        input_schema={"type": "object", "properties": {"test": {"type": "string"}}},
+        handler=example_handler,
+    )
+
+    # Test getting the registered tool
+    tool = registry.get_tool("test_tool")
+    assert tool is not None
+    assert tool.name == "test_tool"
+    assert tool.description == "A test tool"
+    assert callable(tool.handler)
+
+    # Test getting non-existent tool
+    assert registry.get_tool("non_existent") is None
+
+
+def test_list_tools():
+    registry = MCPToolRegistry()
+
+    # Register multiple tools
+    registry.register_tool(
+        name="tool1", description="Tool 1", input_schema={}, handler=example_handler
+    )
+    registry.register_tool(
+        name="tool2", description="Tool 2", input_schema={}, handler=example_handler
+    )
+
+    # Test listing tools
+    tools = registry.list_tools()
+    assert len(tools) == 2
+    assert {tool.name for tool in tools} == {"tool1", "tool2"}
+
+
+def test_load_tools_from_config():
+    registry = MCPToolRegistry()
+
+    # Test valid config
+    valid_config = [
+        {
+            "name": "config_tool",
+            "description": "A tool from config",
+            "input_schema": {"type": "object"},
+            "handler": "test_tool_registry.example_handler",
+        }
+    ]
+
+    registry.load_tools_from_config(valid_config)
+    assert "config_tool" in registry.tools
+    assert registry.tools["config_tool"].name == "config_tool"
+    assert registry.tools["config_tool"].description == "A tool from config"
+    assert callable(registry.tools["config_tool"].handler)
+
+
+def test_tool_execution():
+    registry = MCPToolRegistry()
+
+    # Register a tool
+    registry.register_tool(
+        name="echo",
+        description="Echo the input",
+        input_schema={"type": "object"},
+        handler=example_handler,
+    )
+
+    # Get and execute the tool
+    tool = registry.get_tool("echo")
+    assert tool is not None
+
+    test_input = {"message": "hello"}
+    result = tool.handler(test_input)
+    assert result == {"result": test_input}