From 922f4a533d97bf1aed9b4bb59146411a48286006 Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Mon, 19 May 2025 13:34:35 +0200 Subject: [PATCH] flaky test: handle second possible result in anthropic streaming case --- .../test_anthropic_with_tool_call.py | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/tests/integration/anthropic/test_anthropic_with_tool_call.py b/tests/integration/anthropic/test_anthropic_with_tool_call.py index bbfd893..82cdc82 100644 --- a/tests/integration/anthropic/test_anthropic_with_tool_call.py +++ b/tests/integration/anthropic/test_anthropic_with_tool_call.py @@ -248,14 +248,27 @@ async def test_streaming_response_with_tool_call( messages = [{"role": "user", "content": query}] response = weather_agent.get_streaming_response(messages) - assert response is not None - assert response[0][0].type == "text" - assert response[0][1].type == "tool_use" - assert response[0][1].name == "get_weather" - assert city in response[0][1].input["location"].lower() + if len(response) == 2: + assert response is not None + assert response[0][0].type == "text" + assert response[0][1].type == "tool_use" + assert response[0][1].name == "get_weather" + assert city in response[0][1].input["location"].lower() - assert response[1][0].type == "text" - assert city in response[1][0].text.lower() + assert response[1][0].type == "text" + assert city in response[1][0].text.lower() + elif len(response) == 1: + # expected output in this case is something like this: + # [[TextBlock(text="I'll help you check the weather in New York using the get_weather function.", type='text', citations=None), ToolUseBlock(id='toolu_019VZsmxuUhShou2EpPBxvpe', input={'location': 'New York, NY', 'unit': 'celsius'}, name='get_weather', type='tool_use')]] + + assert response is not None + assert response[0][0].type == "text" + assert response[0][1].type == "tool_use" + assert response[0][1].name == "get_weather" + assert city in response[0][1].input["location"].lower() + + else: + assert False, "Expected response length 2 or 1, but got" + str(response) if push_to_explorer: # Wait for the trace to be saved