diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 40fea7dd7..77ece41fb 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -2,6 +2,7 @@ import sys, os import traceback import pytest from dotenv import load_dotenv +import openai load_dotenv() @@ -56,6 +57,25 @@ def test_openai_azure_embedding_simple(): test_openai_azure_embedding_simple() + +def test_openai_azure_embedding_timeouts(): + try: + response = embedding( + model="azure/azure-embedding-model", + input=["good morning from litellm"], + timeout=0.00001 + ) + print(response) + response_keys = dict(response).keys() + assert set(["usage", "model", "object", "data"]) == set(response_keys) #assert litellm response has expected keys from OpenAI embedding response + except openai.APITimeoutError: + print("Good job got timeout error!") + pass + except Exception as e: + pytest.fail(f"Expected timeout error, did not get the correct error. Instead got {e}") + +# test_openai_azure_embedding_timeouts() + def test_openai_azure_embedding(): try: api_key = os.environ['AZURE_API_KEY']