From f68a6b6e135b5d6c92828df446d7bc26ece26be9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 16 Apr 2024 21:46:24 -0700 Subject: [PATCH] test(test_amazing_vertex_completion.py): handle api errors --- litellm/tests/test_amazing_vertex_completion.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 44201db8317..a7296f03eb2 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -602,6 +602,8 @@ def test_gemini_pro_function_calling(): print(f"completion: {completion}") assert completion.choices[0].message.content is None assert len(completion.choices[0].message.tool_calls) == 1 + except litellm.APIError as e: + pass except litellm.RateLimitError as e: pass except Exception as e: @@ -656,6 +658,8 @@ def test_gemini_pro_function_calling_streaming(): # assert len(completion.choices[0].message.tool_calls) == 1 for chunk in completion: print(f"chunk: {chunk}") + except litellm.APIError as e: + pass except litellm.RateLimitError as e: pass @@ -699,6 +703,8 @@ async def test_gemini_pro_async_function_calling(): print(f"completion: {completion}") assert completion.choices[0].message.content is None assert len(completion.choices[0].message.tool_calls) == 1 + except litellm.APIError as e: + pass except litellm.RateLimitError as e: pass except Exception as e: