mcp-dbutils

"""Integration tests for prompts functionality"""

import asyncio
import tempfile

import anyio
import mcp.types as types
import pytest
import yaml
from mcp import ClientSession

from mcp_dbutils.base import ConnectionServer
from mcp_dbutils.log import create_logger

# 创建测试用的 logger
logger = create_logger("test-prompts", True)  # debug=True 以显示所有日志

@pytest.mark.asyncio
async def test_prompts_capability(sqlite_db, mcp_config):
    """Test that prompts capability is properly set"""
    with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
        yaml.dump(mcp_config, tmp)
        tmp.flush()
        server = ConnectionServer(config_path=tmp.name)

        # Get initialization options and verify prompts capability
        init_options = server.server.create_initialization_options()
        assert init_options.capabilities is not None
        assert init_options.capabilities.prompts is not None

@pytest.mark.asyncio
async def test_list_prompts(sqlite_db, mcp_config):
    """Test that list_prompts returns an empty list"""
    with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as tmp:
        yaml.dump(mcp_config, tmp)
        tmp.flush()
        server = ConnectionServer(config_path=tmp.name)

        # Create bidirectional streams with proper types
        # Client -> Server stream (server receives messages and exceptions)
        client_to_server_send, client_to_server_recv = anyio.create_memory_object_stream[types.JSONRPCMessage | Exception](10)

        # Server -> Client stream (client receives messages)
        server_to_client_send, server_to_client_recv = anyio.create_memory_object_stream[types.JSONRPCMessage](10)

        # Start server in background task with proper stream connections
        server_task = asyncio.create_task(
            server.server.run(
                client_to_server_recv,  # Server reads from client
                server_to_client_send,  # Server writes to client
                server.server.create_initialization_options(),
                raise_exceptions=True  # 让错误更容易调试
            )
        )

        try:
            try:
                # Initialize client with proper stream connections
                logger("debug", "Creating client session")
                client = ClientSession(
                    server_to_client_recv,   # Client reads from server
                    client_to_server_send    # Client writes to server
                )

                async with client:
                    # Initialize the session
                    logger("debug", "Initializing client session")
                    init_response = await client.initialize()
                    logger("debug", f"Client session initialized with response: {init_response}")

                    # Test prompts list with timeout
                    logger("debug", "Sending prompts/list request")
                    try:
                        # 使用 SDK 的 list_prompts 方法
                        response = await asyncio.wait_for(client.list_prompts(), timeout=3.0)
                        logger("debug", f"Got response: {response}")

                        # 验证响应
                        assert isinstance(response.prompts, list)
                        assert len(response.prompts) == 0
                        logger("debug", "Test completed successfully")
                    except asyncio.TimeoutError:
                        logger("error", "Request timed out after 3 seconds")
                        raise
                    except Exception as e:
                        logger("error", f"Request failed: {str(e)}")
                        logger("error", f"Request error type: {type(e)}")
                        raise

            except Exception as e:
                logger("error", f"Test failed with error: {str(e)}")
                logger("error", f"Error type: {type(e)}")
                raise

        except asyncio.TimeoutError:
            logger("error", "Test timed out")
            raise RuntimeError("Test timed out after 3 seconds")

        finally:
            # Clean up
            logger("debug", "Starting cleanup")
            server_task.cancel()
            try:
                await server_task
            except asyncio.CancelledError:
                logger("debug", "Server task cancelled")
            except Exception as e:
                logger("error", f"Error during server cleanup: {str(e)}")

            # Close streams
            try:
                await client_to_server_send.aclose()
                await client_to_server_recv.aclose()
                await server_to_client_send.aclose()
                await server_to_client_recv.aclose()
                logger("debug", "All streams closed")
            except Exception as e:
                logger("error", f"Error during stream cleanup: {str(e)}")