From d1a045f4f4cf6af5aa4219de6086a59c4ab6cabd Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Tue, 29 Sep 2026 11:51:52 -0700
Subject: [PATCH] test(mcp): avoid coverage import overhead in conformance
 peers

---
 tests/integration/_support/mcp.py             |  6 +++--
 tests/integration/mcp/test_mcp_management.py  |  2 +-
 .../integration_support/test_conformance.py   | 23 +++++++++++++------
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py
index 5736e750912..996bd054d7f 100644
--- a/tests/integration/_support/mcp.py
+++ b/tests/integration/_support/mcp.py
@@ -222,8 +222,10 @@ def stdio_peer(directory: Path, *, rich: bool = False) -> Iterator[McpPeer]:
     coverage_file: Final = os.environ.get("COVERAGE_FILE")
     coverage_args: Final = (
         (
-            "-m",
-            "coverage",
+            "-c",
+            "import mcp.server.mcpserver, runpy, sys; "
+            "sys.argv = ['coverage', *sys.argv[1:]]; "
+            "runpy.run_module('coverage', run_name='__main__')",
             "run",
             f"--rcfile={STDIO_PEER.parent.parent / 'conformance_coverage.toml'}",
             f"--data-file={Path(coverage_file).resolve()}",
diff --git a/tests/integration/mcp/test_mcp_management.py b/tests/integration/mcp/test_mcp_management.py
index 67cdbbff5a4..0ffed8fdaac 100644
--- a/tests/integration/mcp/test_mcp_management.py
+++ b/tests/integration/mcp/test_mcp_management.py
@@ -22,7 +22,7 @@ from integration._support.mcp import (
 from integration._support.process import owned_proxy
 
 from litellm.models.user import LiteLLM_UserTable
-from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
+from litellm.proxy.proxy_server import ExperimentalUIJWTToken
 
 ADD: Final = {"a": 4, "b": 5}
 
diff --git a/tests/unit/integration_support/test_conformance.py b/tests/unit/integration_support/test_conformance.py
index 32b1e2ed8cb..4ad56699135 100644
--- a/tests/unit/integration_support/test_conformance.py
+++ b/tests/unit/integration_support/test_conformance.py
@@ -53,16 +53,19 @@ async def test_negotiation_records_preserve_the_peer_call_contract(monkeypatch:
 
 
 @pytest.mark.parametrize("covered", (False, True))
+@pytest.mark.parametrize("has_data_file", (False, True))
 def test_stdio_coverage_preserves_default_environment(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, covered: bool
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, covered: bool, has_data_file: bool
 ) -> None:
     monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2]))
     from integration._support.mcp import STDIO_PEER, stdio_peer
 
     monkeypatch.delenv("COVERAGE_PROCESS_CONFIG", raising=False)
+    monkeypatch.delenv("COVERAGE_FILE", raising=False)
     monkeypatch.setenv("OPENAI_API_KEY", "must-not-reach-peer")
-    coverage_file: Final = tmp_path / "coverage"
-    monkeypatch.setenv("COVERAGE_FILE", str(coverage_file))
+    coverage_file: Final = tmp_path / "coverage data"
+    if has_data_file:
+        monkeypatch.setenv("COVERAGE_FILE", str(coverage_file))
     if covered:
         monkeypatch.setenv("COVERAGE_PROCESS_CONFIG", "coverage-config")
     with stdio_peer(tmp_path) as peer:
@@ -72,13 +75,15 @@ def test_stdio_coverage_preserves_default_environment(
             "args": [
                 *(
                     [
-                        "-m",
-                        "coverage",
+                        "-c",
+                        "import mcp.server.mcpserver, runpy, sys; "
+                        "sys.argv = ['coverage', *sys.argv[1:]]; "
+                        "runpy.run_module('coverage', run_name='__main__')",
                         "run",
                         f"--rcfile={STDIO_PEER.parent.parent / 'conformance_coverage.toml'}",
                         f"--data-file={coverage_file}",
                     ]
-                    if covered
+                    if covered and has_data_file
                     else []
                 ),
                 str(STDIO_PEER),
@@ -96,7 +101,11 @@ def test_peer_drain_preserves_received_requests(
     from integration._support.mcp import McpPeer
 
     request: Final = {"body": {"method": "tools/call"}, "headers": {"authorization": "synthetic"}}
-    negotiation: Final = {"body": {}, "headers": {}, "negotiation": {"requested": "2025-03-26", "returned": "2025-03-26"}}
+    negotiation: Final = {
+        "body": {},
+        "headers": {},
+        "negotiation": {"requested": "2025-03-26", "returned": "2025-03-26"},
+    }
     records: Final = (request, negotiation)
     observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
     path: Final = tmp_path / "peer.jsonl" if file_backed else None