From 376fd86e525328421b45d09c8980168e99c0e9ba Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Tue, 29 Sep 2026 09:31:47 -0700
Subject: [PATCH] test(mcp): preserve stdio environment during coverage

---
 tests/integration/_support/mcp.py             | 25 +++++++++++--------
 .../integration_support/test_conformance.py   | 24 +++++++++++++++---
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py
index 2d88a77a037..5736e750912 100644
--- a/tests/integration/_support/mcp.py
+++ b/tests/integration/_support/mcp.py
@@ -58,16 +58,7 @@ class McpPeer:
 
     def registration(self) -> dict[str, object]:
         if self.transport == "stdio":
-            return {
-                "transport": "stdio",
-                "command": self.command,
-                "args": list(self.args),
-                **(
-                    {"env": {"COVERAGE_PROCESS_CONFIG": os.environ["COVERAGE_PROCESS_CONFIG"]}}
-                    if os.environ.get("COVERAGE_PROCESS_CONFIG")
-                    else {}
-                ),
-            }
+            return {"transport": "stdio", "command": self.command, "args": list(self.args)}
         if self.spec_path is not None:
             return {"transport": "http", "url": self.url, "spec_path": str(self.spec_path)}
         return {"transport": self.transport, "url": self.url}
@@ -228,12 +219,24 @@ def mcp_peer(transport: Literal["http", "sse"] = "http", *, rich: bool = False)
 @contextmanager
 def stdio_peer(directory: Path, *, rich: bool = False) -> Iterator[McpPeer]:
     record: Final = directory / f"stdio-{os.getpid()}-{time.monotonic_ns()}.jsonl"
+    coverage_file: Final = os.environ.get("COVERAGE_FILE")
+    coverage_args: Final = (
+        (
+            "-m",
+            "coverage",
+            "run",
+            f"--rcfile={STDIO_PEER.parent.parent / 'conformance_coverage.toml'}",
+            f"--data-file={Path(coverage_file).resolve()}",
+        )
+        if coverage_file and os.environ.get("COVERAGE_PROCESS_CONFIG")
+        else ()
+    )
     yield McpPeer(
         "",
         queue.Queue(),
         "stdio",
         sys.executable,
-        (str(STDIO_PEER), str(record), "rich" if rich else "plain"),
+        (*coverage_args, str(STDIO_PEER), str(record), "rich" if rich else "plain"),
         record,
     )
 
diff --git a/tests/unit/integration_support/test_conformance.py b/tests/unit/integration_support/test_conformance.py
index aaa8bfa07b1..32b1e2ed8cb 100644
--- a/tests/unit/integration_support/test_conformance.py
+++ b/tests/unit/integration_support/test_conformance.py
@@ -53,22 +53,38 @@ async def test_negotiation_records_preserve_the_peer_call_contract(monkeypatch:
 
 
 @pytest.mark.parametrize("covered", (False, True))
-def test_stdio_registration_forwards_only_explicit_coverage(
+def test_stdio_coverage_preserves_default_environment(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch, covered: bool
 ) -> None:
     monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
-    from integration._support.mcp import stdio_peer
+    from integration._support.mcp import STDIO_PEER, stdio_peer
 
     monkeypatch.delenv("COVERAGE_PROCESS_CONFIG", raising=False)
     monkeypatch.setenv("OPENAI_API_KEY", "must-not-reach-peer")
+    coverage_file: Final = tmp_path / "coverage"
+    monkeypatch.setenv("COVERAGE_FILE", str(coverage_file))
     if covered:
         monkeypatch.setenv("COVERAGE_PROCESS_CONFIG", "coverage-config")
     with stdio_peer(tmp_path) as peer:
         assert peer.registration() == {
             "transport": "stdio",
             "command": peer.command,
-            "args": list(peer.args),
-            **({"env": {"COVERAGE_PROCESS_CONFIG": "coverage-config"}} if covered else {}),
+            "args": [
+                *(
+                    [
+                        "-m",
+                        "coverage",
+                        "run",
+                        f"--rcfile={STDIO_PEER.parent.parent / 'conformance_coverage.toml'}",
+                        f"--data-file={coverage_file}",
+                    ]
+                    if covered
+                    else []
+                ),
+                str(STDIO_PEER),
+                str(peer.record),
+                "plain",
+            ],
         }