From 24395492aa5a231b20d0074825aeec6feb8b41b4 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 30 Jul 2024 21:47:52 -0700
Subject: [PATCH] test: cleanup duplicate tests + add error handling for
 backend api errors

---
 litellm/tests/test_cohere_completion.py | 38 +++++++------------------
 litellm/tests/test_completion.py        |  2 ++
 litellm/tests/test_streaming.py         |  8 ++++--
 3 files changed, 17 insertions(+), 31 deletions(-)

diff --git a/litellm/tests/test_cohere_completion.py b/litellm/tests/test_cohere_completion.py
index 372c87b4000..e90818fee5d 100644
--- a/litellm/tests/test_cohere_completion.py
+++ b/litellm/tests/test_cohere_completion.py
@@ -1,19 +1,23 @@
-import sys, os
+import os
+import sys
 import traceback
+
 from dotenv import load_dotenv
 
 load_dotenv()
-import os, io
+import io
+import os
 
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
-import pytest
-import litellm
-from litellm import embedding, completion, completion_cost, Timeout
-from litellm import RateLimitError
 import json
 
+import pytest
+
+import litellm
+from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
+
 litellm.num_retries = 3
 
 
@@ -37,28 +41,6 @@ def test_chat_completion_cohere():
         pytest.fail(f"Error occurred: {e}")
 
 
-def test_chat_completion_cohere_stream():
-    try:
-        litellm.set_verbose = False
-        messages = [
-            {
-                "role": "user",
-                "content": "Hey",
-            },
-        ]
-        response = completion(
-            model="cohere_chat/command-r",
-            messages=messages,
-            max_tokens=10,
-            stream=True,
-        )
-        print(response)
-        for chunk in response:
-            print(chunk)
-    except Exception as e:
-        pytest.fail(f"Error occurred: {e}")
-
-
 def test_chat_completion_cohere_tool_calling():
     try:
         litellm.set_verbose = True
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 1e180de290e..a01cb0483a4 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -3630,6 +3630,8 @@ def test_chat_completion_cohere_stream():
         print(response)
         for chunk in response:
             print(chunk)
+    except litellm.APIConnectionError as e:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index d2eb08cd333..d708c7155a3 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -567,11 +567,13 @@ async def test_completion_predibase_streaming(sync_mode):
                 raise Exception("Empty response received")
 
         print(f"complete_response: {complete_response}")
-    except litellm.Timeout as e:
+    except litellm.Timeout:
         pass
-    except litellm.InternalServerError as e:
+    except litellm.InternalServerError:
         pass
-    except litellm.ServiceUnavailableError as e:
+    except litellm.ServiceUnavailableError:
+        pass
+    except litellm.APIConnectionError:
         pass
     except Exception as e:
         print("ERROR class", e.__class__)