diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py index 73eb73ab214..e6ae70f27e9 100644 --- a/tests/integration/_support/mcp.py +++ b/tests/integration/_support/mcp.py @@ -44,13 +44,17 @@ class McpPeer: spec_path: Path | None = None consumed: list[int] = field(default_factory=lambda: [0]) - def drain(self) -> tuple[dict[str, object], ...]: + def drain(self, *, include_negotiation: bool = False) -> tuple[dict[str, object], ...]: if self.record is not None: lines: Final = self.record.read_text().splitlines() if self.record.exists() else [] fresh: Final = tuple(json.loads(line) for line in lines[self.consumed[0] :]) self.consumed[0] = len(lines) - return fresh - return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize())) + return tuple(item for item in fresh if include_negotiation or "negotiation" not in item) + return tuple( + item + for item in (self.calls.get_nowait() for _ in range(self.calls.qsize())) + if include_negotiation or "negotiation" not in item + ) def registration(self) -> dict[str, object]: if self.transport == "stdio": diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index 0cd36a9f808..8d2ddcc1e73 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -231,7 +231,7 @@ def test_pinned_revision_pairs_list_and_call_through_gateway( peer.drain() negotiated: Final = asyncio.run(exercise()) - observed: Final = peer.drain() + observed: Final = peer.drain(include_negotiation=True) negotiations: Final = tuple( (item["negotiation"]["requested"], item["negotiation"]["returned"]) for item in observed diff --git a/tests/unit/integration_support/test_conformance.py b/tests/unit/integration_support/test_conformance.py index db102b9cfea..abc8ea8728f 100644 --- a/tests/unit/integration_support/test_conformance.py +++ b/tests/unit/integration_support/test_conformance.py @@ -1,6 +1,7 @@ import hashlib import json import os +import queue from pathlib import Path from typing import Final from unittest.mock import Mock, patch @@ -50,6 +51,33 @@ async def test_negotiation_records_preserve_the_peer_call_contract() -> None: assert records[0]["negotiation"] == {"requested": "2025-03-26", "returned": "2025-03-26"} +@pytest.mark.parametrize("file_backed", (False, True)) +def test_peer_drain_preserves_received_requests(tmp_path: Path, file_backed: bool) -> None: + 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"}} + records: Final = (request, negotiation) + observed: Final[queue.Queue[dict[str, object]]] = queue.Queue() + path: Final = tmp_path / "peer.jsonl" if file_backed else None + if path is not None: + path.write_text("".join(json.dumps(record) + "\n" for record in records)) + else: + for record in records: + observed.put(record) + peer: Final = McpPeer("http://peer", observed, record=path) + assert peer.drain() == (request,) + assert peer.drain() == () + if path is not None: + with path.open("a") as sink: + sink.write("".join(json.dumps(record) + "\n" for record in records)) + else: + for record in records: + observed.put(record) + assert peer.drain(include_negotiation=True) == records + assert peer.drain(include_negotiation=True) == () + + @pytest.mark.parametrize("explicit", (False, True)) def test_owned_proxy_isolates_automatic_coverage_unless_requested( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, explicit: bool