From e6f7b21a8a379e387fcecdf75bbeec63d48b5b72 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 21 Feb 2024 21:45:33 -0800 Subject: [PATCH] test: handle api errors for gemini/palm --- litellm/tests/test_completion.py | 4 ++++ litellm/tests/test_streaming.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 798da53f7f..7260c243cf 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1986,6 +1986,8 @@ def test_completion_gemini(): response = completion(model=model_name, messages=messages) # Add any assertions here to check the response print(response) + except litellm.APIError as e: + pass except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -2015,6 +2017,8 @@ def test_completion_palm(): response = completion(model=model_name, messages=messages) # Add any assertions here to check the response print(response) + except litellm.APIError as e: + pass except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index bf0da99d98..f1640d97da 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -392,6 +392,8 @@ def test_completion_palm_stream(): if complete_response.strip() == "": raise Exception("Empty response received") print(f"completion_response: {complete_response}") + except litellm.APIError as e: + pass except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -425,6 +427,8 @@ def test_completion_gemini_stream(): if complete_response.strip() == "": raise Exception("Empty response received") print(f"completion_response: {complete_response}") + except litellm.APIError as e: + pass except Exception as e: pytest.fail(f"Error occurred: {e}")