DataOps-platform/tests/integration/test_dataops_mcp_container.py

import os
from pathlib import Path

import anyio
import pytest
from mcp import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client

from app.core.mcp.servers import CONTEXT_TOOL_NAMES, SCHEDULING_TOOL_NAMES

pytestmark = pytest.mark.integration
ROOT = Path(__file__).resolve().parents[2]
COMPOSE = ROOT / "deploy/docker/docker-compose.yml"


@pytest.mark.parametrize(
    ("service", "expected"),
    [
        ("dataops-context-mcp", CONTEXT_TOOL_NAMES),
        ("dataops-scheduling-mcp", SCHEDULING_TOOL_NAMES),
    ],
)
def test_built_dataops_mcp_containers_complete_stdio_handshake(service, expected):
    if os.environ.get("RUN_V53_MCP_CONTAINER") != "1":
        pytest.skip("set RUN_V53_MCP_CONTAINER=1 for container acceptance")
    environment = {
        **os.environ,
        "DATAOPS_MCP_SUBJECT": "v53-container-agent",
        "DATAOPS_MCP_ROLES": "scheduler",
        "DATAOPS_MCP_BUSINESS_DOMAINS": "sales",
        "DATAOPS_MCP_ENVIRONMENTS": "test",
        "DATAOPS_MCP_CORRELATION_ID": ("01900000-0000-7000-8000-000000000053"),
    }
    async def list_tools():
        parameters = StdioServerParameters(
            command="docker",
            args=[
                "compose",
                "-f",
                str(COMPOSE),
                "--profile",
                "dataops-mcp",
                "run",
                "--rm",
                "-T",
                "--no-deps",
                service,
            ],
            cwd=str(ROOT),
            env=environment,
        )
        with anyio.fail_after(30):
            async with (
                stdio_client(parameters) as (read_stream, write_stream),
                ClientSession(read_stream, write_stream) as session,
            ):
                await session.initialize()
                response = await session.list_tools()
                return {tool.name for tool in response.tools}

    assert anyio.run(list_tools) == expected