"""Tests for MiddlewareMCPServer implementation."""
import json
from typing import Dict, Any
import pytest
from mcp_chain import MiddlewareMCPServer, mcp_chain
class MockMCPServer:
"""Mock MCP server for testing."""
def __init__(self, metadata: Dict[str, Any], response_data: Dict[str, Any]):
self._metadata = metadata
self._response_data = response_data
def get_metadata(self) -> str:
return json.dumps(self._metadata)
def handle_request(self, request: str) -> str:
return json.dumps(self._response_data)
def test_middleware_basic_proxy():
"""Test that middleware can proxy to downstream server."""
metadata = {"tools": [{"name": "test_tool"}]}
response = {"result": "test_response"}
downstream = MockMCPServer(metadata, response)
middleware = MiddlewareMCPServer(downstream_server=downstream)
# Should proxy metadata
assert json.loads(middleware.get_metadata()) == metadata
# Should proxy requests
request = '{"method": "test"}'
assert json.loads(middleware.handle_request(request)) == response
def test_middleware_no_downstream_error():
"""Test that middleware raises error when no downstream server."""
# MiddlewareMCPServer now requires downstream in constructor
with pytest.raises(TypeError):
MiddlewareMCPServer()
def test_middleware_dict_metadata_transformer():
"""Test middleware with dict metadata transformer."""
# Create a mock dict-based downstream server
class MockDictServer:
def get_metadata(self):
return {"tools": [{"name": "test_tool", "description": "original"}]}
def handle_request(self, request):
return {"result": "success"}
downstream = MockDictServer()
def transform_metadata(next_server, metadata_dict):
# Get metadata from next_server
original_metadata = next_server.get_metadata()
modified = original_metadata.copy()
for tool in modified.get("tools", []):
tool["description"] = "transformed: " + tool.get("description", "")
return modified
middleware = MiddlewareMCPServer(
downstream_server=downstream, metadata_transformer=transform_metadata
)
result = middleware.get_metadata()
expected = {
"tools": [{"name": "test_tool", "description": "transformed: original"}]
}
assert result == expected
def test_middleware_dict_request_transformer():
"""Test middleware with dict request transformer."""
# Create a mock dict-based downstream server
class MockDictServer:
def get_metadata(self):
return {"tools": []}
def handle_request(self, request):
return {"result": "original", "received": request}
downstream = MockDictServer()
def transform_request(next_server, request_dict):
# Transform request
modified_request = request_dict.copy()
modified_request["transformed"] = True
# Get response from next_server
response = next_server.handle_request(modified_request)
# Transform response
modified_response = response.copy()
modified_response["response_transformed"] = True
return modified_response
middleware = MiddlewareMCPServer(
downstream_server=downstream, request_transformer=transform_request
)
# The downstream should receive the transformed request
# and we should get back the transformed response
result = middleware.handle_request({"method": "test"})
assert result["result"] == "original"
assert result["response_transformed"] is True
assert result["received"]["transformed"] is True
assert result["received"]["method"] == "test"
def test_then_with_downstream_server():
"""Test .then() method with downstream server."""
# Create a dict-based downstream server
class MockDictServer:
def get_metadata(self):
return {"tools": [{"name": "test_tool"}]}
def handle_request(self, request):
return {"result": "success"}
downstream = MockDictServer()
# Use mcp_chain() to create the chain and add downstream server
from mcp_chain import mcp_chain
chained = mcp_chain().then(downstream)
# Should return the downstream server directly (since no transformers)
assert chained == downstream
assert chained.get_metadata() == {"tools": [{"name": "test_tool"}]}
def test_then_with_porcelain_metadata_transformer():
"""Test .then() method with porcelain metadata transformer."""
# Create a dict-based downstream server
class MockDictServer:
def get_metadata(self):
return {"tools": [{"name": "test_tool", "description": "original"}]}
def handle_request(self, request):
return {"result": "success"}
downstream = MockDictServer()
def transform_metadata(next_server, metadata_dict: dict) -> dict:
# Get metadata from downstream and transform it
original = next_server.get_metadata()
result = original.copy()
for tool in result.get("tools", []):
tool["description"] = "transformed: " + tool.get("description", "")
return result
def identity_request_transformer(next_server, request_dict: dict) -> dict:
# Simple passthrough for request
return next_server.handle_request(request_dict)
# Use mcp_chain() to create the chain with transformers and downstream
from mcp_chain import mcp_chain, MiddlewareMCPServer
chained = (
mcp_chain()
.then(transform_metadata, identity_request_transformer)
.then(downstream)
)
assert isinstance(chained, MiddlewareMCPServer)
result = chained.get_metadata()
expected = {
"tools": [{"name": "test_tool", "description": "transformed: original"}]
}
assert result == expected
def test_then_with_porcelain_request_transformer():
"""Test .then() method with porcelain request transformer."""
# Create a dict-based downstream server
class MockDictServer:
def get_metadata(self):
return {"tools": []}
def handle_request(self, request):
return {"result": "original", "received": request}
downstream = MockDictServer()
def identity_metadata_transformer(next_server, metadata_dict: dict) -> dict:
# Simple passthrough for metadata
return next_server.get_metadata()
def transform_request(next_server, request_dict: dict) -> dict:
# Transform request and call downstream
modified_request = {**request_dict, "transformed": True}
return next_server.handle_request(modified_request)
# Use mcp_chain() to create the chain with transformers and downstream
from mcp_chain import mcp_chain, MiddlewareMCPServer
chained = (
mcp_chain()
.then(identity_metadata_transformer, transform_request)
.then(downstream)
)
assert isinstance(chained, MiddlewareMCPServer)
result = chained.handle_request({"method": "test"})
# Should show that the request was transformed before reaching downstream
assert result["result"] == "original"
assert result["received"]["method"] == "test"
assert result["received"]["transformed"] is True
def test_then_with_raw_transformers():
"""Test .then() method with raw transformers."""
metadata = {"tools": [{"name": "test_tool", "description": "original"}]}
downstream = MockMCPServer(metadata, {"result": "original"})
def raw_metadata_transformer(next_mcp, json_metadata: str) -> str:
# Get metadata from next_mcp
original_metadata = next_mcp.get_metadata()
data = json.loads(original_metadata)
for tool in data.get("tools", []):
tool["description"] = "raw_transformed: " + tool.get("description", "")
return json.dumps(data)
def raw_request_transformer(next_mcp, json_request: str) -> str:
request = json.loads(json_request)
request["raw_transformed"] = True
# Get response from next_mcp
actual_response = next_mcp.handle_request(json.dumps(request))
response = json.loads(actual_response)
response["raw_response_transformed"] = True
return json.dumps(response)
# Use mcp_chain() to create the chain with transformers and downstream
from mcp_chain import mcp_chain, MiddlewareMCPServer
chained = (
mcp_chain()
.then(raw_metadata_transformer, raw_request_transformer)
.then(downstream)
)
assert isinstance(chained, MiddlewareMCPServer)
# Test metadata transformation
metadata_result = json.loads(chained.get_metadata())
expected_metadata = {
"tools": [{"name": "test_tool", "description": "raw_transformed: original"}]
}
assert metadata_result == expected_metadata
# Test request transformation
request_result = json.loads(chained.handle_request('{"method": "test"}'))
expected_response = {"result": "original", "raw_response_transformed": True}
assert request_result == expected_response
def test_factory_function():
"""Test mcp_chain factory function."""
from mcp_chain import MCPChainBuilder
chain = mcp_chain()
assert isinstance(chain, MCPChainBuilder)
def test_chaining_multiple_middleware():
"""Test chaining multiple middleware together with correct execution order."""
# Create downstream server with initial data
metadata = {"tools": [{"name": "test_tool", "description": "downstream"}]}
response = {"result": "downstream"}
downstream = MockMCPServer(metadata, response)
def first_metadata_transformer(next_mcp, json_metadata: str) -> str:
"""First transformer - should execute first in the request flow."""
original_metadata = next_mcp.get_metadata()
metadata = json.loads(original_metadata)
result = metadata.copy()
for tool in result.get("tools", []):
tool["description"] = "first: " + tool.get("description", "")
return json.dumps(result)
def first_request_transformer(next_mcp, json_request: str) -> str:
"""First transformer - should process requests first."""
request = json.loads(json_request)
# Transform request
modified_request = {**request, "first_processed": True}
# Get response from downstream
response = next_mcp.handle_request(json.dumps(modified_request))
response_data = json.loads(response)
# Transform response on the way back
transformed_response = {**response_data, "first_response": True}
return json.dumps(transformed_response)
def second_metadata_transformer(next_mcp, json_metadata: str) -> str:
"""Second transformer - should execute second in the request flow."""
original_metadata = next_mcp.get_metadata()
metadata = json.loads(original_metadata)
result = metadata.copy()
for tool in result.get("tools", []):
tool["description"] = "second: " + tool.get("description", "")
return json.dumps(result)
def second_request_transformer(next_mcp, json_request: str) -> str:
"""Second transformer - should process requests second."""
request = json.loads(json_request)
# Transform request
modified_request = {**request, "second_processed": True}
# Get response from downstream
response = next_mcp.handle_request(json.dumps(modified_request))
response_data = json.loads(response)
# Transform response on the way back
transformed_response = {**response_data, "second_response": True}
return json.dumps(transformed_response)
# Build chain: client → first → second → downstream
chain = (
mcp_chain()
.then(first_metadata_transformer, first_request_transformer)
.then(second_metadata_transformer, second_request_transformer)
.then(downstream)
)
# Test metadata transformation order
# Expected flow: first calls second calls downstream
# second transforms "downstream" → "second: downstream"
# first transforms "second: downstream" → "first: second: downstream"
metadata_result = json.loads(chain.get_metadata())
expected_metadata = {
"tools": [{"name": "test_tool", "description": "first: second: downstream"}]
}
assert metadata_result == expected_metadata
# Test request transformation order
# Expected flow:
# 1. second processes request: adds "second_processed": True
# 2. first processes request: adds "first_processed": True
# 3. downstream receives: {"method": "test", "second_processed": True, "first_processed": True}
# 4. downstream responds: {"result": "downstream"}
# 5. first processes response: adds "first_response": True
# 6. second processes response: adds "second_response": True
request_result = json.loads(chain.handle_request('{"method": "test"}'))
expected_response = {
"result": "downstream",
"first_response": True,
"second_response": True,
}
assert request_result == expected_response
