From 1b215d704db2acec5a388652d35fbf207863d4bc Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 17 Jun 2024 16:33:58 -0700
Subject: [PATCH] test: cleanup tests

---
 litellm/tests/test_completion.py |  2 ++
 litellm/tests/test_streaming.py  | 24 +++++++++++++++---------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 7010fa4f48..79e859f5cf 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -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"
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 622b2efc84..ecb21b9f2b 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -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):