test: cleanup tests

This commit is contained in:
Krrish Dholakia 2024-06-17 16:33:58 -07:00
parent 57d9479252
commit 1b215d704d
2 changed files with 17 additions and 9 deletions

View file

@ -3459,6 +3459,7 @@ def test_completion_palm_stream():
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="Account deleted by IBM.")
def test_completion_watsonx():
litellm.set_verbose = True
model_name = "watsonx/ibm/granite-13b-chat-v2"
@ -3479,6 +3480,7 @@ def test_completion_watsonx():
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="Skip test. account deleted.")
def test_completion_stream_watsonx():
litellm.set_verbose = True
model_name = "watsonx/ibm/granite-13b-chat-v2"

View file

@ -1,12 +1,17 @@
#### What this tests ####
# This tests streaming for the completion endpoint
import sys, os, asyncio
import asyncio
import os
import sys
import time
import traceback
import time, pytest, uuid
from pydantic import BaseModel
import uuid
from typing import Tuple
import pytest
from pydantic import BaseModel
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
@ -15,12 +20,12 @@ from dotenv import load_dotenv
load_dotenv()
import litellm
from litellm import (
completion,
acompletion,
AuthenticationError,
BadRequestError,
RateLimitError,
ModelResponse,
RateLimitError,
acompletion,
completion,
)
litellm.logging = False
@ -1644,9 +1649,8 @@ def test_sagemaker_weird_response():
When the stream ends, flush any remaining holding chunks.
"""
try:
from litellm.llms.sagemaker import TokenIterator
import json
import json
from litellm.llms.sagemaker import TokenIterator
chunk = """<s>[INST] Hey, how's it going? [/INST],
@ -1772,6 +1776,7 @@ def test_completion_sagemaker_stream():
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="Account deleted by IBM.")
def test_completion_watsonx_stream():
litellm.set_verbose = True
try:
@ -2631,9 +2636,10 @@ def test_success_callback_streaming():
# test_success_callback_streaming()
from typing import List, Optional
#### STREAMING + FUNCTION CALLING ###
from pydantic import BaseModel
from typing import List, Optional
class Function(BaseModel):