mcp-dbutils
by donghao1393
- tests
- integration
"""Extended SQLite handler integration tests"""
import os
import tempfile
import pytest
import pytest_asyncio
import yaml
from mcp_dbutils.base import (
ConnectionHandlerError,
ConnectionServer,
)
from mcp_dbutils.log import create_logger
# 创建测试用的 logger
logger = create_logger("test-sqlite-extended", True) # debug=True 以显示所有日志
@pytest_asyncio.fixture(autouse=True)
async def setup_test_table(sqlite_db, mcp_config):
"""Create test table before each test"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
# Create test table
await handler.execute_query("""
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY,
name TEXT,
email TEXT
)
""")
# Insert test data
await handler.execute_query("""
INSERT OR IGNORE INTO users (id, name, email) VALUES
(1, 'Test User 1', 'test1@example.com'),
(2, 'Test User 2', 'test2@example.com')
""")
yield
# Cleanup is handled by the sqlite_db fixture
@pytest.mark.asyncio
async def test_get_table_indexes(sqlite_db, mcp_config):
"""Test getting index information for SQLite table"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
# Get indexes for users table
indexes = await handler.get_table_indexes("users")
# Verify indexes content
assert "No indexes found on table users" in indexes
# SQLite might not have indexes for the test table by default
# so we're just checking the header is present
@pytest.mark.asyncio
async def test_get_table_indexes_nonexistent(sqlite_db, mcp_config):
"""Test getting indexes for nonexistent table"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
with pytest.raises(ConnectionHandlerError):
await handler.get_table_indexes("nonexistent_table")
@pytest.mark.asyncio
async def test_explain_query(sqlite_db, mcp_config):
"""Test the query explanation functionality"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
# Get execution plan for a SELECT query
explain_result = await handler.explain_query("SELECT * FROM users WHERE id = 1")
# Verify the explanation includes expected SQLite EXPLAIN output
assert "Query Execution Plan:" in explain_result
assert "Details:" in explain_result
@pytest.mark.asyncio
async def test_explain_query_invalid(sqlite_db, mcp_config):
"""Test the query explanation with invalid query"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
with pytest.raises(ConnectionHandlerError):
await handler.explain_query("SELECT * FROM nonexistent_table")
@pytest.mark.asyncio
async def test_get_table_stats(sqlite_db, mcp_config):
"""Test getting table statistics"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
# Get statistics for users table
stats = await handler.get_table_stats("users")
# Verify statistics content
assert "Table Statistics for users:" in stats
assert "Row Count:" in stats
assert "Total Size:" in stats
@pytest.mark.asyncio
async def test_get_table_stats_nonexistent(sqlite_db, mcp_config):
"""Test getting statistics for nonexistent table"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
with pytest.raises(ConnectionHandlerError):
await handler.get_table_stats("nonexistent_table")
@pytest.mark.asyncio
async def test_execute_complex_queries(sqlite_db, mcp_config):
"""Test executing more complex SELECT queries"""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
yaml.dump(mcp_config, tmp)
tmp.flush()
server = ConnectionServer(config_path=tmp.name)
async with server.get_handler("test_sqlite") as handler:
# Query with GROUP BY
result_str = await handler.execute_query(
"SELECT COUNT(*) as count FROM users GROUP BY name"
)
result = eval(result_str)
assert "columns" in result
assert "rows" in result
assert len(result["columns"]) == 1
assert result["columns"][0] == "count"
# Query with ORDER BY and LIMIT
result_str = await handler.execute_query(
"SELECT name FROM users ORDER BY name LIMIT 1"
)
result = eval(result_str)
assert len(result["rows"]) == 1
# Query with JOIN (assuming a related table exists, otherwise this may fail)
try:
# Create a posts table if it doesn't exist
path = mcp_config["connections"]["test_sqlite"]["path"]
os.system(f"""
sqlite3 {path} '
CREATE TABLE IF NOT EXISTS posts (
id INTEGER PRIMARY KEY,
user_id INTEGER,
title TEXT,
FOREIGN KEY (user_id) REFERENCES users(id)
);
INSERT OR IGNORE INTO posts (id, user_id, title) VALUES (1, 1, "Test Post 1");
INSERT OR IGNORE INTO posts (id, user_id, title) VALUES (2, 2, "Test Post 2");
'
""")
# Execute JOIN query
result_str = await handler.execute_query(
"SELECT u.name, p.title FROM users u JOIN posts p ON u.id = p.user_id ORDER BY u.name"
)
result = eval(result_str)
assert len(result["columns"]) == 2
assert "name" in result["columns"]
assert "title" in result["columns"]
except Exception as e:
logger.warning(f"JOIN query test skipped: {str(e)}")