conftest.py

#
#  Copyright (C) 2017-2025 Dremio Corporation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

"""
Global pytest fixtures for dremio-mcp tests.
"""
import os
import random
import uuid
from typing import AsyncGenerator, NamedTuple

import pytest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import patch
from collections import OrderedDict

from dremioai.config import settings
from dremioai.config.tools import ToolType
from dremioai.servers.mcp import (
    Transports,
    init,
    create_metrics_server,
)

from mocks.http_mock import (
    create_pytest_logging_server_fixture,
    start_server_with_app,
    ServerFixture,
    LoggingServerFixture,
)
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
import contextlib
from dremioai.log import set_level
from dremioai.metrics import registry
from prometheus_client import CollectorRegistry


@pytest.fixture(autouse=True)
def reset_uvicorn_logger_propagation():
    """Reset uvicorn logger state between tests.

    Uvicorn's configure_logging() sets uvicorn.access.propagate=False via its
    default LOGGING_CONFIG when a server starts. This leaks into subsequent tests
    that assert stdlib loggers propagate to the root handler.
    """
    yield
    import logging

    for name in ("uvicorn.access", "uvicorn.error", "uvicorn"):
        lg = logging.getLogger(name)
        lg.propagate = True
        for h in lg.handlers[:]:
            lg.removeHandler(h)


@pytest.fixture(autouse=True)
def reset_sse_starlette_app_status():
    """
    Reset the global AppStatus.should_exit_event from sse_starlette between tests.

    This fixes the asyncio event loop binding issue where the global event gets
    bound to the first event loop and causes "bound to a different event loop"
    errors in subsequent tests.
    """
    try:
        # Import and reset the global state
        from sse_starlette.sse import AppStatus

        AppStatus.should_exit_event = None
    except ImportError:
        # sse_starlette might not be available in all test environments
        pass

    yield

    # Clean up after test
    try:
        from sse_starlette.sse import AppStatus

        AppStatus.should_exit_event = None
    except ImportError:
        pass


@pytest.fixture(autouse=True)
def reset_metrics_registry():
    """
    Reset the global metrics registry between tests.

    This ensures that metrics from previous tests don't persist and affect subsequent
    test assertions.
    """
    registry._registry = CollectorRegistry()

    yield


@pytest.fixture
def temp_config_dir():
    """Create a temporary directory for config files"""
    with TemporaryDirectory() as temp_dir:
        yield Path(temp_dir)


@pytest.fixture
def mock_config_dir(temp_config_dir):
    """Mock the home directory to use our temporary directory"""
    with patch.object(Path, "home", return_value=temp_config_dir):
        # Also patch XDG_CONFIG_HOME environment variable
        old_env = os.environ.get("XDG_CONFIG_HOME")
        os.environ["XDG_CONFIG_HOME"] = str(temp_config_dir)
        yield temp_config_dir
        # Restore original environment
        if old_env:
            os.environ["XDG_CONFIG_HOME"] = old_env
        else:
            os.environ.pop("XDG_CONFIG_HOME", None)


@pytest.fixture
def mock_settings_instance():
    """Create a mock settings instance with default values"""
    old_settings = settings.instance()
    try:
        settings._settings.set(
            settings.Settings.model_validate(
                {
                    "dremio": {
                        "uri": "https://test-dremio-uri.com",
                        "pat": "test-pat",
                        "project_id": uuid.uuid4(),
                    },
                    "tools": {"server_mode": ToolType.FOR_SELF.name},
                }
            )
        )
        yield settings.instance()
    finally:
        settings._settings.set(old_settings)


@pytest.fixture
def temp_config_dir():
    """Create a temporary directory for config files"""
    with TemporaryDirectory() as temp_dir:
        yield Path(temp_dir)


@pytest.fixture
def mock_config_dir(temp_config_dir):
    """Mock the home directory to use our temporary directory"""
    with patch.object(Path, "home", return_value=temp_config_dir):
        # Also patch XDG_CONFIG_HOME environment variable
        old_env = os.environ.get("XDG_CONFIG_HOME")
        os.environ["XDG_CONFIG_HOME"] = str(temp_config_dir)
        yield temp_config_dir
        # Restore original environment
        if old_env:
            os.environ["XDG_CONFIG_HOME"] = old_env
        else:
            os.environ.pop("XDG_CONFIG_HOME", None)


def _create_logging_server(log_level="warning"):
    # Mock data for HTTP endpoints that tools will call
    mock_data = OrderedDict(
        [
            (r"/sql", "sql/job_submission.json"),  # SQL query submission
            (r"/job/test-job-12345$", "sql/job_status.json"),  # Job status check
            (r"/job/test-job-12345/results$", "sql/job_results.json"),  # Job results
            (r"/search", "search/search_results.json"),  # Search endpoints
            (r"/catalog/.*/wiki", "catalog/wiki.json"),  # Wiki endpoints
            (r"/catalog/.*/tags", "catalog/tags.json"),  # Tags endpoints
            (r"/catalog/.*/graph", "catalog/lineage.json"),  # Lineage endpoints
            (r"/catalog(/by-path)?", "catalog/table_schema.json"),  # Schema endpoints
        ]
    )

    return create_pytest_logging_server_fixture(
        mock_data=mock_data, port=8000, log_level=log_level
    )


@pytest.fixture
def logging_level(request: pytest.FixtureRequest):
    if request.config.get_verbosity() > 2:
        return "debug"
    if request.config.get_verbosity() > 1:
        return "info"
    return "warning"


@pytest.fixture
def logging_server(logging_level):
    server = _create_logging_server(logging_level)
    try:
        yield server
    finally:
        server.close()


class StreamableMcpServerFixture(NamedTuple):
    mcp_server: ServerFixture
    logging_server: LoggingServerFixture
    metrics_port: int


@contextlib.asynccontextmanager
async def http_streamable_mcp_server(
    logging_server: LoggingServerFixture,
    logging_level: str,
    project_id: str = None,
    wlm_engine: str = None,
    dremio_overrides: dict = None,
) -> AsyncGenerator[StreamableMcpServerFixture, None]:
    old = settings.instance()
    sf = None
    try:
        settings.configure(force=True)
        host = "127.0.0.1"
        port = random.randrange(9000, 12000)
        metrics_port = random.randrange(9000, 12000)

        # Ensure metrics port is different from main port
        while metrics_port == port:
            metrics_port = random.randrange(9000, 12000)

        config = {
            "dremio": {
                "uri": logging_server.url,
                "project_id": uuid.uuid4(),
                "pat": "test-pat",
                "enable_search": True,
                "metrics_enabled": True,
                "metrics": {
                    "enabled": True,
                    "port": metrics_port,
                },
            },
            "tools": {"server_mode": ToolType.FOR_DATA_PATTERNS.name},
        }
        if wlm_engine:
            config["dremio"]["wlm"] = {"engine_name": wlm_engine}
        if dremio_overrides:
            config["dremio"].update(dremio_overrides)
        settings._settings.set(settings.Settings.model_validate(config))
        settings.write_settings()

        set_level(logging_level.upper())

        # Start metrics server using asyncio
        metrics_server = create_metrics_server(
            host=host, port=metrics_port, log_level=logging_level
        )

        mcp_server = init(
            transport=Transports.streamable_http,
            port=port,
            mode=settings.instance().tools.server_mode,
            support_project_id_endpoints=project_id is not None,
        )

        app = mcp_server.streamable_http_app()
        server, stop_event = start_server_with_app(
            app,
            host=host,
            port=port,
            log_level=logging_level,
            additional_runners=[metrics_server.serve()],
        )
        sf = ServerFixture(
            f"http://{host}:{port}/mcp/{(str(project_id) + '/') if project_id else ''}",
            stop_event,
            server,
        )

        yield StreamableMcpServerFixture(sf, logging_server, metrics_port)
    finally:
        if sf is not None:
            sf.close()
        print(f"{sf} closed")
        settings._settings.set(old)


@contextlib.asynccontextmanager
async def http_streamable_client_server(
    sf: ServerFixture, token=None
) -> AsyncGenerator[ClientSession, None]:
    headers = {"Authorization": f"Bearer {token}"} if token is not None else None
    async with streamablehttp_client(url=sf.url, headers=headers) as (
        read_stream,
        write_stream,
        gid,
    ):
        print(f"Client connected to {sf.url}")
        async with ClientSession(read_stream, write_stream) as session:
            await session.initialize()
            yield session