From e948b39e3ab545b6502e9159d9415d6c55f91a66 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 25 Jan 2024 18:34:13 -0800 Subject: [PATCH] test(test_streaming.py): fix test to handle none chunk --- litellm/tests/test_streaming.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index d9f99bece75..fda640c962f 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -847,9 +847,13 @@ def test_sagemaker_weird_response(): logging_obj=logging_obj, ) complete_response = "" - for chunk in response: - print(chunk) - complete_response += chunk["choices"][0]["delta"]["content"] + for idx, chunk in enumerate(response): + # print + chunk, finished = streaming_format_tests(idx, chunk) + has_finish_reason = finished + complete_response += chunk + if finished: + break assert len(complete_response) > 0 except Exception as e: pytest.fail(f"An exception occurred - {str(e)}")