Tavily MCP Server

"""Test helpers for extracting handlers from the server module."""
import asyncio
import inspect
from unittest.mock import patch, MagicMock, AsyncMock
import mcp_server_tavily.server as server_module


class ServerHandlerExtractor:
    """Extract handler functions from the server module."""
    
    def __init__(self):
        self.handlers = {}
        
    async def extract_handlers(self):
        """Extract all handler functions from the server module."""
        # Create mocks for all decorators
        with patch('mcp_server_tavily.server.Server') as mock_server_class:
            mock_server = MagicMock()
            mock_server_class.return_value = mock_server
            
            # Mock the decorator methods to capture handlers
            def capture_handler(decorator_name):
                def decorator(func):
                    self.handlers[decorator_name] = func
                    return func
                return decorator
            
            mock_server.list_tools.side_effect = capture_handler('list_tools')
            mock_server.list_prompts.side_effect = capture_handler('list_prompts')
            mock_server.call_tool.side_effect = capture_handler('call_tool')
            mock_server.get_prompt.side_effect = capture_handler('get_prompt')
            
            # Mock stdio_server to prevent actual I/O
            with patch('mcp_server_tavily.server.stdio_server') as mock_stdio:
                mock_stdio.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock())
                
                # Mock TavilyClient to prevent actual API calls
                with patch('mcp_server_tavily.server.TavilyClient') as mock_client_class:
                    # Call serve to register all handlers
                    await server_module.serve("fake_api_key")
        
        return self.handlers


async def get_server_handlers():
    """Get all handler functions from the server module."""
    extractor = ServerHandlerExtractor()
    return await extractor.extract_handlers()