Selaa lähdekoodia

test: stabilize DataOps MCP container handshake

马小龙 4 päivää sitten
vanhempi
commit
0a992913b6
1 muutettua tiedostoa jossa 31 lisäystä ja 58 poistoa
  1. 31 58
      tests/integration/test_dataops_mcp_container.py

+ 31 - 58
tests/integration/test_dataops_mcp_container.py

@@ -1,9 +1,10 @@
-import json
 import os
-import subprocess
 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
 
@@ -12,30 +13,6 @@ 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"),
     [
@@ -54,35 +31,31 @@ def test_built_dataops_mcp_containers_complete_stdio_handshake(service, expected
         "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
+    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