fix clarifai - test

This commit is contained in:
Ishaan Jaff 2024-05-11 16:54:22 -07:00
parent 2e9f4ff23a
commit f0c727a597

View file

@ -11,7 +11,15 @@ sys.path.insert(
) # Adds the parent directory to the system path ) # Adds the parent directory to the system path
import pytest import pytest
import litellm import litellm
from litellm import embedding, completion, acompletion, acreate, completion_cost, Timeout, ModelResponse from litellm import (
embedding,
completion,
acompletion,
acreate,
completion_cost,
Timeout,
ModelResponse,
)
from litellm import RateLimitError from litellm import RateLimitError
# litellm.num_retries = 3 # litellm.num_retries = 3
@ -20,6 +28,7 @@ litellm.success_callback = []
user_message = "Write a short poem about the sky" user_message = "Write a short poem about the sky"
messages = [{"content": user_message, "role": "user"}] messages = [{"content": user_message, "role": "user"}]
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True)
def reset_callbacks(): def reset_callbacks():
print("\npytest fixture - resetting callbacks") print("\npytest fixture - resetting callbacks")
@ -28,6 +37,7 @@ def reset_callbacks():
litellm.failure_callback = [] litellm.failure_callback = []
litellm.callbacks = [] litellm.callbacks = []
def test_completion_clarifai_claude_2_1(): def test_completion_clarifai_claude_2_1():
print("calling clarifai claude completion") print("calling clarifai claude completion")
import os import os
@ -67,6 +77,7 @@ def test_completion_clarifai_mistral_large():
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio @pytest.mark.asyncio
def test_async_completion_clarifai(): def test_async_completion_clarifai():
import asyncio import asyncio
@ -89,5 +100,4 @@ def test_async_completion_clarifai():
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred: {e}") pytest.fail(f"An exception occurred: {e}")
asyncio.run(test_get_response()) asyncio.run(test_get_response())