StarTree MCP Server for Apache Pinot

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

from mcp.server import NotificationOptions, Server
import mcp.types as types
import pytest

from mcp_pinot.server import create_server, run_http_server, run_stdio_server


class TestCreateServer:
    """Test the create_server function"""

    def test_create_server_returns_server_instance(self):
        """Test that create_server returns a Server instance"""
        server = create_server()
        assert isinstance(server, Server)
        assert server.name == "pinot_mcp"

    @pytest.mark.asyncio
    async def test_create_server_handlers_work(self):
        """Test that the server handlers are properly registered"""
        server = create_server()

        # Test that the server has the expected capabilities
        capabilities = server.get_capabilities(
            notification_options=NotificationOptions(), experimental_capabilities={}
        )
        assert capabilities.prompts is not None
        assert capabilities.tools is not None

    @pytest.mark.asyncio
    async def test_create_server_tools_handler(self):
        """Test that the tools handler is properly registered"""
        server = create_server()

        # Test that the server has tools capability
        capabilities = server.get_capabilities(
            notification_options=NotificationOptions(), experimental_capabilities={}
        )
        assert capabilities.tools is not None

        # Test that we can create the server without errors
        # (The actual tool registration is tested in the existing server tests)
        assert server.name == "pinot_mcp"


class TestHttpTransport:
    """Test HTTP transport functionality"""

    @pytest.mark.asyncio
    async def test_run_stdio_server(self):
        """Test the STDIO server runner"""
        with patch("mcp_pinot.server.mcp.server.stdio.stdio_server") as mock_stdio:
            mock_read_stream = AsyncMock()
            mock_write_stream = AsyncMock()
            mock_stdio.return_value.__aenter__.return_value = (
                mock_read_stream,
                mock_write_stream,
            )

            # Mock the server.run method to avoid actually running
            with patch("mcp_pinot.server.create_server") as mock_create_server:
                mock_server = MagicMock()
                mock_server.run = AsyncMock()
                mock_server.get_capabilities.return_value = types.ServerCapabilities(
                    supportsPrompts=True,
                    supportsTools=True,
                    supportsNotifications=True,
                    supportsExperimentalCapabilities=True,
                )
                mock_create_server.return_value = mock_server

                await run_stdio_server()

                # Verify server.run was called with correct parameters
                mock_server.run.assert_called_once()
                args = mock_server.run.call_args[0]
                assert args[0] == mock_read_stream
                assert args[1] == mock_write_stream

    @pytest.mark.asyncio
    async def test_run_http_server_basic(self):
        """Test the HTTP server runner with basic configuration"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.host = "127.0.0.1"
            mock_server_config.port = 8080
            mock_server_config.endpoint = "/sse"
            mock_server_config.ssl_keyfile = None
            mock_server_config.ssl_certfile = None

            with patch("mcp_pinot.server.uvicorn.Server") as mock_uvicorn_server:
                mock_server_instance = AsyncMock()
                mock_uvicorn_server.return_value = mock_server_instance

                # Create a task that we can cancel to avoid infinite running
                task = asyncio.create_task(run_http_server())
                await asyncio.sleep(0.01)  # Let it start
                task.cancel()

                try:
                    await task
                except asyncio.CancelledError:
                    pass  # Expected

                # Verify uvicorn server was created
                mock_uvicorn_server.assert_called_once()

    @pytest.mark.asyncio
    async def test_run_http_server_with_ssl(self):
        """Test the HTTP server runner with SSL configuration"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.host = "0.0.0.0"
            mock_server_config.port = 8443
            mock_server_config.endpoint = "/sse"
            mock_server_config.ssl_keyfile = "/path/to/key.pem"
            mock_server_config.ssl_certfile = "/path/to/cert.pem"

            with patch("mcp_pinot.server.ssl.SSLContext") as mock_ssl_context:
                with patch("mcp_pinot.server.uvicorn.Server") as mock_uvicorn_server:
                    mock_server_instance = AsyncMock()
                    mock_uvicorn_server.return_value = mock_server_instance

                    # Create a task that we can cancel
                    task = asyncio.create_task(run_http_server())
                    await asyncio.sleep(0.01)  # Let it start
                    task.cancel()

                    try:
                        await task
                    except asyncio.CancelledError:
                        pass  # Expected

                    # Verify SSL context was created
                    mock_ssl_context.assert_called_once()


class TestHttpEndpoints:
    """Test HTTP endpoint handling"""

    @pytest.fixture
    def mock_scope(self):
        """Create a mock ASGI scope for HTTP requests"""
        return {
            "type": "http",
            "method": "GET",
            "path": "/sse",
            "query_string": b"",
            "headers": [],
        }

    @pytest.fixture
    def mock_receive(self):
        """Create a mock ASGI receive callable"""
        return AsyncMock()

    @pytest.fixture
    def mock_send(self):
        """Create a mock ASGI send callable"""
        return AsyncMock()

    @pytest.mark.asyncio
    async def test_http_app_sse_get_request(self, mock_scope, mock_receive, mock_send):
        """Test that GET requests to SSE endpoint are handled correctly"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.endpoint = "/sse"

            with patch(
                "mcp_pinot.server.mcp.server.sse.SseServerTransport"
            ) as mock_transport:
                mock_transport_instance = MagicMock()
                mock_transport.return_value = mock_transport_instance

                # Mock the connect_sse context manager
                mock_streams = (AsyncMock(), AsyncMock())
                mock_transport_instance.connect_sse.return_value.__aenter__ = AsyncMock(
                    return_value=mock_streams
                )
                mock_transport_instance.connect_sse.return_value.__aexit__ = AsyncMock()

                with patch("mcp_pinot.server.create_server") as mock_create_server:
                    mock_server = MagicMock()
                    mock_server.run = AsyncMock()
                    mock_server.get_capabilities.return_value = (
                        types.ServerCapabilities(
                            supportsPrompts=True,
                            supportsTools=True,
                            supportsNotifications=True,
                            supportsExperimentalCapabilities=True,
                        )
                    )
                    mock_create_server.return_value = mock_server

                    # Import the app function from run_http_server
                    # We need to extract the app function for testing
                    # This is a bit tricky since it's defined inside run_http_server
                    # For now, we'll test the overall flow
                    pass

    @pytest.mark.asyncio
    async def test_http_app_post_request(self, mock_receive, mock_send):
        """Test that POST requests to SSE endpoint are handled correctly"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.endpoint = "/sse"

            with patch(
                "mcp_pinot.server.mcp.server.sse.SseServerTransport"
            ) as mock_transport:
                mock_transport_instance = MagicMock()
                mock_transport.return_value = mock_transport_instance
                mock_transport_instance.handle_post_message = AsyncMock()

                # This test would require extracting the ASGI app from run_http_server
                # For now, we verify the mocking setup
                pass

    @pytest.mark.asyncio
    async def test_http_app_404_response(self, mock_receive, mock_send):
        """Test that requests to unknown paths return 404"""
        # This would test the 404 handling in the ASGI app
        # Implementation would require extracting the app function
        pass


class TestMainFunction:
    """Test the main function with transport selection"""

    @pytest.mark.asyncio
    async def test_main_with_stdio_transport(self):
        """Test main function selects STDIO transport"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.transport = "stdio"

            with patch("mcp_pinot.server.run_stdio_server") as mock_run_stdio:
                mock_run_stdio.return_value = AsyncMock()

                from mcp_pinot.server import main

                await main()

                mock_run_stdio.assert_called_once()

    @pytest.mark.asyncio
    async def test_main_with_http_transport(self):
        """Test main function selects HTTP transport"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.transport = "http"

            with patch("mcp_pinot.server.run_http_server") as mock_run_http:
                mock_run_http.return_value = AsyncMock()

                from mcp_pinot.server import main

                await main()

                mock_run_http.assert_called_once()

    @pytest.mark.asyncio
    async def test_main_with_invalid_transport(self):
        """Test main function raises error for invalid transport"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.transport = "invalid"

            from mcp_pinot.server import main

            with pytest.raises(ValueError, match="Unknown transport"):
                await main()

    @pytest.mark.asyncio
    async def test_main_with_both_transports(self):
        """Test main function runs both transports simultaneously"""
        with patch("mcp_pinot.server.server_config") as mock_server_config:
            mock_server_config.transport = "both"

            with patch("mcp_pinot.server.run_stdio_server") as mock_run_stdio:
                with patch("mcp_pinot.server.run_http_server") as mock_run_http:
                    # Make both return immediately to avoid hanging
                    mock_run_stdio.return_value = None
                    mock_run_http.return_value = None

                    # Mock asyncio.wait to return immediately
                    with patch("asyncio.wait") as mock_wait:
                        # Simulate both tasks completing successfully
                        mock_task1 = MagicMock()
                        mock_task1.exception.return_value = None
                        mock_task2 = MagicMock()
                        mock_task2.exception.return_value = None

                        mock_wait.return_value = ([mock_task1, mock_task2], [])

                        from mcp_pinot.server import main

                        await main()

                        # Verify both transports were started
                        mock_run_stdio.assert_called_once()
                        mock_run_http.assert_called_once()


class TestIntegration:
    """Integration tests for HTTP transport"""

    @pytest.mark.asyncio
    async def test_server_creation_and_config_integration(self):
        """Test that server creation works with different configurations"""
        # Test with default config
        with patch.dict("os.environ", {}, clear=True):
            from mcp_pinot.config import load_server_config

            config = load_server_config()
            assert config.transport == "both"

            server = create_server()
            assert server.name == "pinot_mcp"

        # Test with HTTP config
        with patch.dict("os.environ", {"MCP_TRANSPORT": "http"}, clear=True):
            config = load_server_config()
            assert config.transport == "http"

            server = create_server()
            assert server.name == "pinot_mcp"

    @pytest.mark.asyncio
    async def test_transport_selection_integration(self):
        """Test that transport selection works end-to-end"""
        # Test STDIO selection
        with patch("mcp_pinot.server.server_config") as mock_config:
            mock_config.transport = "stdio"

            with patch("mcp_pinot.server.run_stdio_server") as mock_stdio:
                with patch("mcp_pinot.server.run_http_server") as mock_http:
                    from mcp_pinot.server import main

                    await main()

                    mock_stdio.assert_called_once()
                    mock_http.assert_not_called()

        # Test HTTP selection
        with patch("mcp_pinot.server.server_config") as mock_config:
            mock_config.transport = "http"

            with patch("mcp_pinot.server.run_stdio_server") as mock_stdio:
                with patch("mcp_pinot.server.run_http_server") as mock_http:
                    from mcp_pinot.server import main

                    await main()

                    mock_http.assert_called_once()
                    mock_stdio.assert_not_called()