DataOps-platform/tests/integration/test_dataops_mcp_container.py

import json
import os
import subprocess
from pathlib import Path

import pytest

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"


def _messages():
    return [
        {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2025-06-18",
                "capabilities": {},
                "clientInfo": {
                    "name": "dataops-v53-container-test",
                    "version": "1.0",
                },
            },
        },
        {
            "jsonrpc": "2.0",
            "method": "notifications/initialized",
            "params": {},
        },
        {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}},
    ]


@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"),
    }
    result = subprocess.run(
        [
            "docker",
            "compose",
            "-f",
            str(COMPOSE),
            "--profile",
            "dataops-mcp",
            "run",
            "--rm",
            "-T",
            "--no-deps",
            service,
        ],
        cwd=ROOT,
        env=environment,
        input="\n".join(json.dumps(item) for item in _messages()) + "\n",
        capture_output=True,
        text=True,
        timeout=30,
    )

    assert result.returncode == 0, result.stderr
    responses = {
        item["id"]: item
        for line in result.stdout.splitlines()
        if line.strip().startswith("{")
        for item in [json.loads(line)]
        if "id" in item
    }
    assert 2 in responses, (result.stdout, result.stderr)
    assert {tool["name"] for tool in responses[2]["result"]["tools"]} == expected