From 8e430fcfbecb01dd8d6085a2d42bc2fdecd09b7c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 17 Nov 2023 18:01:52 -0800 Subject: [PATCH] test(test_async_fn.py): catch timeouts --- litellm/tests/test_async_fn.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_async_fn.py b/litellm/tests/test_async_fn.py index bd3976ec2..308abe26e 100644 --- a/litellm/tests/test_async_fn.py +++ b/litellm/tests/test_async_fn.py @@ -20,6 +20,8 @@ def test_sync_response(): try: response = completion(model="gpt-3.5-turbo", messages=messages, timeout=5) print(f"response: {response}") + except litellm.Timeout as e: + pass except Exception as e: pytest.fail(f"An exception occurred: {e}") # test_sync_response() @@ -29,12 +31,11 @@ def test_sync_response_anyscale(): user_message = "Hello, how are you?" messages = [{"content": user_message, "role": "user"}] try: - response = completion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages) + response = completion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, timeout=5) + except litellm.Timeout as e: + pass except Exception as e: pytest.fail(f"An exception occurred: {e}") -test_sync_response() -print(f"STARTING ANYSCALE RESPONSE") -test_sync_response_anyscale() # test_sync_response_anyscale() def test_async_response_openai():