import tempfile
import pytest
import yaml
from mcp_dbutils.base import (
ConnectionHandlerError,
ConnectionServer,
)
@pytest.mark.asyncio
async def test_list_tables(sqlite_db, mcp_config):
"""Test listing tables in SQLite database"""
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:
tables = await handler.get_tables()
table_names = [table.name.replace(" schema", "") for table in tables]
assert "products" in table_names
# Check schema information
schema_str = await handler.get_schema("products")
schema = eval(schema_str)
assert schema["columns"][0]["name"] == "id"
assert schema["columns"][0]["type"] == "INTEGER"
assert schema["columns"][1]["name"] == "name"
assert schema["columns"][1]["type"] == "TEXT"
assert schema["columns"][2]["name"] == "price"
assert schema["columns"][2]["type"] == "REAL"
assert schema["columns"][0]["primary_key"] is True
@pytest.mark.asyncio
async def test_execute_query(sqlite_db, mcp_config):
"""Test executing 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:
# Simple SELECT
result_str = await handler.execute_query("SELECT name FROM products ORDER BY price")
result = eval(result_str)
assert len(result["rows"]) == 2
assert result["rows"][0]["name"] == "Widget" # $9.99
assert result["rows"][1]["name"] == "Gadget" # $19.99
# SELECT with WHERE clause
result_str = await handler.execute_query(
"SELECT * FROM products WHERE price < 10.00"
)
result = eval(result_str)
assert len(result["rows"]) == 1
assert result["rows"][0]["name"] == "Widget"
assert float(result["rows"][0]["price"]) == 9.99
@pytest.mark.asyncio
async def test_non_select_query(sqlite_db, mcp_config):
"""Test executing non-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:
# 我们现在允许非SELECT查询
result = await handler.execute_query("DELETE FROM products")
assert result == "Query executed successfully"
@pytest.mark.asyncio
async def test_invalid_query(sqlite_db, mcp_config):
"""Test handling of invalid SQL 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:
with pytest.raises(ConnectionHandlerError):
await handler.execute_query("SELECT * FROM nonexistent_table")
@pytest.mark.asyncio
async def test_connection_cleanup(sqlite_db, mcp_config):
"""Test that database connections are properly cleaned up"""
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:
await handler.get_tables()