"""
Integration tests for MCP protocol compliance.
These tests verify that the MCP server correctly implements the Model Context Protocol
specification, including JSON-RPC format, tools/list, tools/call, and error handling.
"""
import asyncio
import os
import subprocess
import time
from pathlib import Path
from unittest.mock import patch
import pytest
from td_mcp_server.mcp_impl import mcp
class TestMCPProtocolCompliance:
"""Test MCP protocol compliance according to specification."""
async def test_mcp_tools_list_protocol(self):
"""Test tools/list method returns proper MCP protocol response."""
# Get tools using FastMCP's list_tools method
tools = await mcp.list_tools()
# Verify we have the expected number of tools
assert len(tools) == 23, f"Expected 23 tools, got {len(tools)}"
# Verify each tool has required MCP protocol fields
expected_tools = [
# Core database tools
"td_list_databases",
"td_get_database",
"td_list_tables",
# Project management tools
"td_list_projects",
"td_get_project",
"td_download_project_archive",
"td_list_project_files",
"td_read_project_file",
# Workflow tools
"td_list_workflows",
# Search tools
"td_find_project",
"td_find_workflow",
"td_get_project_by_name",
"td_smart_search",
# URL tools
"td_analyze_url",
"td_get_workflow",
# Exploration tools
"td_explore_project",
# Diagnostic tools
"td_diagnose_workflow",
"td_trace_data_lineage",
# Workflow execution tools
"td_get_session",
"td_list_sessions",
"td_get_attempt",
"td_get_attempt_tasks",
"td_analyze_execution",
]
tool_names = [tool.name for tool in tools]
assert set(tool_names) == set(
expected_tools
), f"Tool names mismatch: {tool_names}"
# Verify each tool has required MCP protocol structure
for tool in tools:
# Required fields according to MCP spec
assert hasattr(tool, "name"), f"Tool missing 'name' field: {tool}"
assert hasattr(
tool, "description"
), f"Tool missing 'description' field: {tool}"
assert tool.name, f"Tool name is empty: {tool}"
assert tool.description, f"Tool description is empty: {tool}"
# Verify tool names match expected pattern
assert tool.name in expected_tools, f"Unexpected tool name: {tool.name}"
# Verify description is meaningful (not empty or placeholder)
assert (
len(tool.description) > 10
), f"Tool description too short: {tool.description}"
assert (
"TODO" not in tool.description.upper()
), f"Tool has placeholder description: {tool.description}"
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
async def test_mcp_tool_call_protocol_simple(self, mock_client_class):
"""Test tools/call protocol with simple tool (td_list_databases)."""
from td_mcp_server.api import Database
from td_mcp_server.mcp_impl import td_list_databases
# Setup mock data
mock_databases = [
Database(
name="test_db1",
created_at="2023-01-01 00:00:00 UTC",
updated_at="2023-01-01 00:00:00 UTC",
count=5,
organization=None,
permission="administrator",
delete_protected=False,
),
Database(
name="test_db2",
created_at="2023-01-02 00:00:00 UTC",
updated_at="2023-01-02 00:00:00 UTC",
count=10,
organization=None,
permission="administrator",
delete_protected=True,
),
]
# Setup mock client
mock_client = mock_client_class.return_value
mock_client.get_databases.return_value = mock_databases
# Test with environment variables
with patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
):
# Test default parameters (non-verbose)
result = await td_list_databases()
# Verify MCP protocol compliance for tool response
assert isinstance(
result, dict
), f"Tool should return dict, got {type(result)}"
assert "databases" in result, f"Missing 'databases' key in result: {result}"
assert isinstance(
result["databases"], list
), f"'databases' should be list, got {type(result['databases'])}"
assert result["databases"] == [
"test_db1",
"test_db2",
], f"Unexpected database names: {result['databases']}"
# Test verbose mode
result_verbose = await td_list_databases(verbose=True)
assert isinstance(result_verbose, dict), "Verbose result should be dict"
assert (
"databases" in result_verbose
), "Missing 'databases' key in verbose result"
assert isinstance(
result_verbose["databases"], list
), "Verbose 'databases' should be list"
assert (
len(result_verbose["databases"]) == 2
), "Should have 2 databases in verbose mode"
# Verify verbose mode returns full database objects
for db in result_verbose["databases"]:
assert isinstance(db, dict), "Each database should be a dict"
assert "name" in db, "Database object missing 'name' field"
assert "count" in db, "Database object missing 'count' field"
assert "permission" in db, "Database object missing 'permission' field"
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
async def test_mcp_tool_call_protocol_with_parameters(self, mock_client_class):
"""Test tools/call protocol with parameters (td_list_tables)."""
from td_mcp_server.api import Database, Table
from td_mcp_server.mcp_impl import td_list_tables
# Setup mock data
mock_database = Database(
name="test_db",
created_at="2023-01-01 00:00:00 UTC",
updated_at="2023-01-01 00:00:00 UTC",
count=2,
organization=None,
permission="administrator",
delete_protected=False,
)
mock_tables = [
Table(
id=1,
name="table1",
estimated_storage_size=1000,
counter_updated_at="2023-01-01T00:00:00Z",
last_log_timestamp="2023-01-01T00:00:00Z",
delete_protected=False,
created_at="2023-01-01 00:00:00 UTC",
updated_at="2023-01-01 00:00:00 UTC",
type="log",
include_v=True,
count=100,
table_schema='[["id","string"]]',
expire_days=None,
),
Table(
id=2,
name="table2",
estimated_storage_size=2000,
counter_updated_at="2023-01-02T00:00:00Z",
last_log_timestamp="2023-01-02T00:00:00Z",
delete_protected=True,
created_at="2023-01-02 00:00:00 UTC",
updated_at="2023-01-02 00:00:00 UTC",
type="log",
include_v=True,
count=200,
table_schema='[["id","string"],["value","integer"]]',
expire_days=30,
),
]
# Setup mock client
mock_client = mock_client_class.return_value
mock_client.get_database.return_value = mock_database
mock_client.get_tables.return_value = mock_tables
# Test with environment variables
with patch.dict(
os.environ, {"TD_API_KEY": "test_key", "TD_ENDPOINT": "api.example.com"}
):
# Test required parameter
result = await td_list_tables(database_name="test_db")
# Verify MCP protocol compliance
assert isinstance(result, dict), "Result should be dict"
assert "database" in result, "Missing 'database' key"
assert "tables" in result, "Missing 'tables' key"
assert result["database"] == "test_db", "Database name should match input"
assert isinstance(result["tables"], list), "'tables' should be list"
assert result["tables"] == [
"table1",
"table2",
], "Table names should match mock data"
# Test with pagination parameters
result_paginated = await td_list_tables(
database_name="test_db", limit=10, offset=5, verbose=True
)
# Verify pagination parameters are handled
assert isinstance(result_paginated, dict), "Paginated result should be dict"
assert "tables" in result_paginated, "Paginated result missing 'tables'"
mock_client.get_tables.assert_called_with(
"test_db", limit=10, offset=5, all_results=False
)
@pytest.mark.asyncio
async def test_mcp_error_handling_protocol(self):
"""Test MCP protocol error handling compliance."""
from td_mcp_server.mcp_impl import td_get_database
# Test missing API key
with patch.dict(os.environ, {}, clear=True):
result = await td_get_database(database_name="test_db")
# Verify error response format
assert isinstance(result, dict), "Error response should be dict"
assert "error" in result, "Error response missing 'error' key"
assert isinstance(result["error"], str), "Error message should be string"
assert (
"TD_API_KEY" in result["error"]
), "Error should mention missing API key"
assert (
"environment variable" in result["error"]
), "Error should mention environment variable"
# Test invalid input
with patch.dict(os.environ, {"TD_API_KEY": "test_key"}):
result = await td_get_database(database_name="")
# Verify input validation error
assert isinstance(result, dict), "Validation error should be dict"
assert "error" in result, "Validation error missing 'error' key"
assert (
"cannot be empty" in result["error"]
), "Should validate empty database name"
@pytest.mark.asyncio
@patch("td_mcp_server.mcp_impl.TreasureDataClient")
async def test_mcp_tool_parameter_validation(self, mock_client_class):
"""Test MCP tool parameter validation and type handling."""
from td_mcp_server.mcp_impl import td_list_databases
# Setup mock
mock_client = mock_client_class.return_value
mock_client.get_databases.return_value = []
with patch.dict(os.environ, {"TD_API_KEY": "test_key"}):
# Test default parameters
result = await td_list_databases()
mock_client.get_databases.assert_called_with(
limit=30, offset=0, all_results=False
)
# Test explicit parameters with correct types
result = await td_list_databases(
verbose=True, limit=50, offset=10, all_results=True
)
mock_client.get_databases.assert_called_with(
limit=50, offset=10, all_results=True
)
# Verify boolean parameter handling
assert isinstance(
result, dict
), "Result should be dict even with all parameters"
def test_mcp_server_tool_schema_compliance(self):
"""Test that MCP tools have proper schema definitions."""
# This test verifies that FastMCP can introspect our tools properly
import inspect
from td_mcp_server.diagnostic_tools import (
td_diagnose_workflow,
td_trace_data_lineage,
)
from td_mcp_server.exploration_tools import td_explore_project
from td_mcp_server.mcp_impl import (
td_download_project_archive,
td_get_database,
td_get_project,
td_list_databases,
td_list_project_files,
td_list_projects,
td_list_tables,
td_list_workflows,
td_read_project_file,
)
from td_mcp_server.search_tools import (
td_find_project,
td_find_workflow,
td_get_project_by_name,
td_smart_search,
)
from td_mcp_server.url_tools import td_analyze_url, td_get_workflow
tools = [
# Core database tools
td_list_databases,
td_get_database,
td_list_tables,
# Project management tools
td_list_projects,
td_get_project,
td_download_project_archive,
td_list_project_files,
td_read_project_file,
# Workflow tools
td_list_workflows,
# Search tools
td_find_project,
td_find_workflow,
td_get_project_by_name,
td_smart_search,
# URL tools
td_analyze_url,
td_get_workflow,
# Exploration tools
td_explore_project,
# Diagnostic tools
td_diagnose_workflow,
td_trace_data_lineage,
]
for tool_func in tools:
# Verify function has proper signature
sig = inspect.signature(tool_func)
# Verify function is async
assert asyncio.iscoroutinefunction(
tool_func
), f"{tool_func.__name__} should be async"
# Verify function has docstring
assert tool_func.__doc__, f"{tool_func.__name__} missing docstring"
assert (
len(tool_func.__doc__.strip()) > 20
), f"{tool_func.__name__} docstring too short"
# Verify function has type annotations
for param_name, param in sig.parameters.items():
if param_name != "return":
assert param.annotation != inspect.Parameter.empty, (
f"{tool_func.__name__} parameter '{param_name}' "
"missing type annotation"
)
# Verify return type annotation
assert (
sig.return_annotation != inspect.Parameter.empty
), f"{tool_func.__name__} missing return type annotation"
@pytest.mark.asyncio
async def test_mcp_concurrent_tool_calls(self):
"""Test MCP server handles concurrent tool calls properly."""
from td_mcp_server.mcp_impl import td_list_databases
with patch.dict(os.environ, {"TD_API_KEY": "test_key"}):
with patch(
"td_mcp_server.mcp_impl.TreasureDataClient"
) as mock_client_class:
mock_client = mock_client_class.return_value
mock_client.get_databases.return_value = []
# Run multiple concurrent tool calls
tasks = [
td_list_databases(verbose=False),
td_list_databases(verbose=True),
td_list_databases(limit=10),
td_list_databases(all_results=True),
]
results = await asyncio.gather(*tasks)
# Verify all calls completed successfully
assert len(results) == 4, "All concurrent calls should complete"
for result in results:
assert isinstance(result, dict), "Each result should be dict"
assert (
"databases" in result
), "Each result should have 'databases' key"
class TestMCPServerIntegration:
"""Test MCP server integration at the process level."""
def test_server_startup_with_valid_api_key(self):
"""Test that server.py starts up properly with valid environment."""
env = os.environ.copy()
env["TD_API_KEY"] = "test_key"
env["TD_ENDPOINT"] = "api.example.com"
# Start server process
process = subprocess.Popen(
["uv", "run", "td_mcp_server/server.py"],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=env,
cwd=Path(__file__).parent.parent.parent,
text=True,
)
try:
# Give server time to start
time.sleep(1)
# Check if process is still running (didn't crash on startup)
poll_result = process.poll()
if poll_result is not None:
# Process exited, check stderr for errors
_, stderr = process.communicate()
pytest.fail(
f"Server process exited with code {poll_result}. Stderr: {stderr}"
)
# Server is running successfully
assert process.poll() is None, "Server should be running"
finally:
# Clean up
process.terminate()
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
def test_server_exit_without_api_key(self):
"""Test that server.py exits properly when API key is missing."""
env = os.environ.copy()
# Remove TD_API_KEY if it exists
env.pop("TD_API_KEY", None)
# Start server process
process = subprocess.Popen(
["uv", "run", "td_mcp_server/server.py"],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=env,
cwd=Path(__file__).parent.parent.parent,
text=True,
)
try:
# Wait for process to exit
return_code = process.wait(timeout=10)
# Should exit with error code (1)
assert (
return_code == 1
), f"Server should exit with code 1, got {return_code}"
except subprocess.TimeoutExpired:
process.kill()
process.wait()
pytest.fail("Server should have exited quickly without API key")
def test_server_import_resolution(self):
"""Test that server.py resolves imports correctly in different contexts."""
# This test verifies the import fallback mechanism works
# Test 1: Run as module (should use relative imports)
env = os.environ.copy()
env["TD_API_KEY"] = "test_key"
process = subprocess.Popen(
["uv", "run", "python", "-m", "td_mcp_server.server"],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=env,
cwd=Path(__file__).parent.parent.parent,
text=True,
)
try:
time.sleep(1)
# Should not crash due to import errors
poll_result = process.poll()
if poll_result is not None:
_, stderr = process.communicate()
# Check if it's an import error
if "ImportError" in stderr or "ModuleNotFoundError" in stderr:
pytest.fail(f"Import error when running as module: {stderr}")
finally:
process.terminate()
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
