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

View file

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