From daf22ec8712e406f502da4268efcb1743d5564ff Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 27 Jul 2026 19:22:52 -0700 Subject: [PATCH] test(e2e): make MCP and prometheus e2e tests robust to data-plane sync lag (#34854) * test(e2e): harden harness and tests against data-plane pod churn A stage autoscaler scale-down produced a 2s window of ALB 502s that killed six budget tests on their first management call, and a freshly scaled-up pod that had not run its 30s DB object sync yet failed two MCP tests and one prometheus cardinality test. Retry transient gateway errors (502/503/504, connection errors) once at the shared e2e_http dispatch seam, poll MCP server registration to the poll deadline instead of asserting a single-shot listing, anchor the MCP guardrail full-sync wait to the later of the guardrail and server writes, and turn the prometheus alias poll into a drive-and-scrape convergence loop that re-sends traffic for missing aliases and unions results across scrapes * test(e2e): drain request body in retry stub handler so keep-alive reuse cannot misparse leftovers as requests * revert(e2e): drop the transient-502 retry seam A raw 502 during a pod scale-down is what a real client sees, so the suite retrying past it hides an availability gap instead of flagging it. The gateway-side fix is graceful drain on the deployment; until then the failures are signal * test(e2e): cap per-alias driver re-drives in the prometheus cardinality poll Bounds worst-case provider spend to 4 completions per alias while scrapes keep polling to the deadline; counters persist on whichever pod served them, so the cap costs no convergence unless that pod dies * test(e2e): drop driver re-drives from the prometheus cardinality poll The per-key cardinality contract is process-local and counters persist on whichever pod served the driver call, so unioning aliases across free scrape polls converges without re-sending billable traffic. The residual gap, a pod dying inside the poll window, is deferred to direct per-pod scraping --- .../test_prometheus_cardinality_e2e.py | 12 ++-- tests/e2e/mcp/mcp_client.py | 21 +++++++ tests/e2e/mcp/test_mcp_access_group_e2e.py | 7 +-- tests/e2e/mcp/test_mcp_guardrail_e2e.py | 59 ++++++++++++++----- tests/e2e/mcp/test_mcp_key_access_e2e.py | 9 +-- 5 files changed, 75 insertions(+), 33 deletions(-) diff --git a/tests/e2e/logging/test_prometheus_cardinality_e2e.py b/tests/e2e/logging/test_prometheus_cardinality_e2e.py index 44e3d93c07b..e2d164d8b4c 100644 --- a/tests/e2e/logging/test_prometheus_cardinality_e2e.py +++ b/tests/e2e/logging/test_prometheus_cardinality_e2e.py @@ -9,8 +9,11 @@ stamping ``api_key_alias`` (or collapses every key onto one series) would drop the aliases and fail here. Scraping goes through ``transport.probe`` (raw text) and is parsed with -prometheus_client; the metric is eventually consistent (it increments on the -success-logging callback), so the scrape polls to a deadline. +prometheus_client. ``/metrics`` is per-pod behind a round-robin LB and the metric +is eventually consistent (it increments on the success-logging callback), so the +poll unions the aliases seen across scrapes until the deadline: counters persist +on whichever pod served the driver call, so repeated scrapes converge without +re-sending any billable traffic. """ from __future__ import annotations @@ -58,13 +61,14 @@ class TestPrometheusPerKeyCardinality: deadline = time.monotonic() + client.proxy.poll_timeout seen: frozenset[str] = frozenset() while time.monotonic() < deadline: - seen = _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL) + seen = seen | _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL) if wanted <= seen: break time.sleep(client.proxy.poll_interval) missing = wanted - seen assert not missing, ( - f"{REQUESTS_METRIC} is missing a per-key series for aliases {sorted(missing)}; " + f"{REQUESTS_METRIC} never exposed a per-key series for aliases {sorted(missing)} " + f"on any scraped pod within the deadline; " f"each distinct {ALIAS_LABEL} must grow its own series" ) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 4b1725bb205..6d0f6ddc760 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -195,6 +195,27 @@ class McpClient: ) ).root + def await_registered(self, server_id: str) -> None: + """Poll /v1/mcp/server until `server_id` is listed. Fails at poll_timeout. + + The DB row exists the moment registration returns, but a data-plane pod + answers the listing from a registry it refreshes on a periodic DB sync, so a + pod that joined the load balancer after the write reports the server as + absent until its first sync. + """ + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + registered = frozenset(row.server_id for row in self.registered_servers()) + if server_id in registered: + return + if time.monotonic() >= deadline: + raise AssertionError( + f"registered server {server_id} still absent from /v1/mcp/server " + f"{self.proxy.poll_timeout}s after registration (the data plane never synced " + f"the row): {registered}" + ) + time.sleep(self.proxy.poll_interval) + def generate_key( self, *, diff --git a/tests/e2e/mcp/test_mcp_access_group_e2e.py b/tests/e2e/mcp/test_mcp_access_group_e2e.py index d7ff9736896..1b53d1ca0b4 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/tests/e2e/mcp/test_mcp_access_group_e2e.py @@ -44,12 +44,7 @@ class TestMcpAccessGroupToolSelection: ) resources.defer(lambda: client.proxy.delete_key(other)) - granted_tools = unwrap(client.list_tools(granted)) - assert granted_tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) is not None, ( - f"key granted access group {group} did not see the tagged server's tool " - f"(upstream dead or access-group grant not applied): " - f"{granted_tools.tool_names_for_server(server_id)}" - ) + _ = client.await_tool(granted, server_id, SEARCH_LOGS_TOOL) other_tools = unwrap(client.list_tools(other)).tool_names_for_server(server_id) assert other_tools == frozenset(), ( diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py index dcab235465f..9e3a8c48395 100644 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ b/tests/e2e/mcp/test_mcp_guardrail_e2e.py @@ -29,11 +29,11 @@ from mcp_client import McpCallToolResponse, McpClient, McpToolArguments pytestmark = pytest.mark.e2e # Stage runs several data-plane pods behind the shared key, and each picks up a -# newly registered guardrail only on its next periodic DB sync (~30s in +# newly registered guardrail or MCP server only on its next periodic DB sync (~30s in # proxy_server.py). Every pod is guaranteed to have refreshed only once a full sync -# interval has elapsed since the create; before then a banned call routed to a -# lagging pod passes through as legitimate in-flight propagation, not a leak. -GUARDRAIL_FULL_SYNC_SECONDS = 40.0 +# interval has elapsed since the later of those two writes; before then a banned call +# routed to a lagging pod passes through as legitimate in-flight propagation, not a leak. +FULL_SYNC_SECONDS = 40.0 POST_SYNC_VERIFICATION_CALLS = 4 @@ -53,6 +53,30 @@ def _poll_until_blocked( return last +def _pod_lacks_mcp_server(result: Result[McpCallToolResponse]) -> bool: + """True when the pod that served the call answered as though the MCP server or its + tool does not exist (500 "Tool ... not found"), i.e. its MCP registry has not synced + yet and the request never reached the guardrail at all.""" + if not isinstance(result, UnknownApiError) or result.status_code != 500: + return False + body = result.body.lower() + return "not found" in body and ("tool" in body or "server" in body) + + +def _search_on_synced_pod( + search: Callable[[str], Result[McpCallToolResponse]], query: str, client: McpClient +) -> Result[McpCallToolResponse]: + """Issue `query`, retrying to the poll deadline only while the serving pod does not + know the MCP server yet. Every other outcome, guardrail block or pass-through, comes + back untouched so the caller's assertion still decides it.""" + deadline = time.monotonic() + client.proxy.poll_timeout + last = search(query) + while _pod_lacks_mcp_server(last) and time.monotonic() < deadline: + time.sleep(client.proxy.poll_interval) + last = search(query) + return last + + class TestMcpToolCallGuardrail: @pytest.mark.covers( "guardrail.litellm_content_filter.pre_mcp_call.blocks", @@ -72,6 +96,7 @@ class TestMcpToolCallGuardrail: resources.defer(lambda: client.delete_guardrail(guardrail_id)) server_id = register_datadog_mcp(client, resources) + server_registered_at = time.monotonic() key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id]) resources.defer(lambda: client.proxy.delete_key(key)) @@ -105,31 +130,33 @@ class TestMcpToolCallGuardrail: case _: pytest.fail( "content_filter never blocked the banned keyword on the MCP tool call within " - f"{client.proxy.poll_timeout}s (guardrail sync to the data plane never landed); " - f"last result: {blocked}" + f"{client.proxy.poll_timeout}s (the guardrail or the MCP server never synced to " + f"the data plane); last result: {blocked}" ) # The block above only proves the one pod that served it has synced; another # pod could still lack the guardrail and let the banned call reach Datadog. - # Wait out the full sync interval from the create so every pod has refreshed - # from the DB, then require the banned call to stay blocked across several - # attempts. A pass-through now is a genuine partial-propagation leak, not a - # race. Client load balancing still can't guarantee every pod is hit, so this + # Wait out the full sync interval from the later of the guardrail create and the + # MCP server registration (each syncs on its own clock, so the earlier write's + # deadline can elapse while a pod still lacks the other) so every pod has + # refreshed from the DB, then require the banned call to stay blocked across + # several attempts. A pass-through now is a genuine partial-propagation leak, not + # a race. Client load balancing still can't guarantee every pod is hit, so this # samples several worker selections rather than proving all pods synced. - sync_remaining = guardrail_created_at + GUARDRAIL_FULL_SYNC_SECONDS - time.monotonic() + sync_remaining = max(guardrail_created_at, server_registered_at) + FULL_SYNC_SECONDS - time.monotonic() if sync_remaining > 0: time.sleep(sync_remaining) for attempt in range(1, POST_SYNC_VERIFICATION_CALLS + 1): - reblocked = search(f"still about {banned_keyword} #{attempt}") + reblocked = _search_on_synced_pod(search, f"still about {banned_keyword} #{attempt}", client) assert isinstance(reblocked, UnknownApiError) and reblocked.status_code == 400, ( - "after the guardrail sync interval every data-plane pod must block the banned " - f"keyword, but attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was allowed " - f"through (a pod still lacks the guardrail): {reblocked}" + "after the sync interval every data-plane pod must block the banned keyword, but " + f"attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was not blocked (a pod still " + f"lacks the guardrail, or never synced the MCP server): {reblocked}" ) if attempt < POST_SYNC_VERIFICATION_CALLS: time.sleep(client.proxy.poll_interval) - allowed = search(f"e2e-clean-{marker}") + allowed = _search_on_synced_pod(search, f"e2e-clean-{marker}", client) match allowed: case Success(data=result): assert result.is_error is not True, ( diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 4aeb811a64f..35c864c07d8 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -30,11 +30,6 @@ def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str return key -def _assert_registered(client: McpClient, server_id: str) -> None: - registered = {row.server_id for row in client.registered_servers()} - assert server_id in registered, f"registered server {server_id} absent from /v1/mcp/server: {registered}" - - class TestMcpKeyWithoutAccessIsDenied: @pytest.mark.covers("mcp.list_tools.api_key.denied_without_permission") def test_list_tools_denied_without_permission( @@ -43,7 +38,7 @@ class TestMcpKeyWithoutAccessIsDenied: resources: ResourceManager, ) -> None: server_id = register_datadog_mcp(client, resources) - _assert_registered(client, server_id) + client.await_registered(server_id) permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) @@ -63,7 +58,7 @@ class TestMcpKeyWithoutAccessIsDenied: resources: ResourceManager, ) -> None: server_id = register_datadog_mcp(client, resources) - _assert_registered(client, server_id) + client.await_registered(server_id) permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None)