From 0ffde69391b0a73633e6e737e4cf0d4527d5d494 Mon Sep 17 00:00:00 2001 From: Hemang Sarkar Date: Thu, 15 May 2025 10:54:12 +0200 Subject: [PATCH] Fix test_mcp_sse_with_gateway_hybrid_guardrails so that we don't rely on order of annotations. --- tests/integration/mcp/test_mcp.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/integration/mcp/test_mcp.py b/tests/integration/mcp/test_mcp.py index 3925585..d977509 100644 --- a/tests/integration/mcp/test_mcp.py +++ b/tests/integration/mcp/test_mcp.py @@ -395,15 +395,24 @@ async def test_mcp_sse_with_gateway_hybrid_guardrails( ] # Validate the annotations - annotations = trace["annotations"] + food_annotation = None + tool_call_annotation = None + assert len(annotations) == 2 + for annotation in annotations: + if ( + annotation["content"] == "food in ToolOutput" + and annotation["address"] == "messages.1.content.0.text:22-26" + ): + food_annotation = annotation + elif ( + annotation["content"] == "get_last_message_from_user is called" + and annotation["address"] == "messages.0.tool_calls.0" + ): + tool_call_annotation = annotation + assert food_annotation is not None, "Missing 'food in ToolOutput' annotation" assert ( - annotations[0]["content"] == "get_last_message_from_user is called" - and annotations[0]["address"] == "messages.0.tool_calls.0" - ) - assert annotations[0]["extra_metadata"]["source"] == "guardrails-error" - assert ( - annotations[1]["content"] == "food in ToolOutput" - and annotations[1]["address"] == "messages.1.content.0.text:22-26" - ) - assert annotations[1]["extra_metadata"]["source"] == "guardrails-error" + tool_call_annotation is not None + ), "Missing 'get_last_message_from_user is called' annotation" + assert food_annotation["extra_metadata"]["source"] == "guardrails-error" + assert tool_call_annotation["extra_metadata"]["source"] == "guardrails-error"