From dfdb17ae367d990338641079b3adebf751182853 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 12 Dec 2023 16:21:41 -0800
Subject: [PATCH] test: refactor testing

---
 litellm/tests/test_completion.py | 45 --------------------------------
 litellm/tests/test_streaming.py  | 22 ++++++++++++++++
 2 files changed, 22 insertions(+), 45 deletions(-)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 7b999328626..b8270ba7a11 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -1086,51 +1086,6 @@ def test_completion_chat_sagemaker():
         pytest.fail(f"Error occurred: {e}")
 # test_completion_chat_sagemaker()
 
-
-def test_completion_chat_sagemaker():
-    try:
-        messages = [{"role": "user", "content": "Hey, how's it going?"}]
-        litellm.set_verbose=True
-        response = completion(
-            model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4", 
-            messages=messages,
-            max_tokens=100,
-            temperature=0.7,
-            stream=True,
-        )
-        # Add any assertions here to check the response 
-        complete_response = "" 
-        for chunk in response:
-            complete_response += chunk.choices[0].delta.content or "" 
-        print(f"complete_response: {complete_response}")
-        assert len(complete_response) > 0
-    except Exception as e:
-        pytest.fail(f"Error occurred: {e}")
-
-import asyncio
-@pytest.mark.asyncio
-async def test_completion_chat_sagemaker(): 
-    try: 
-        messages = [{"role": "user", "content": "Hey, how's it going?"}]
-        litellm.set_verbose=True
-        response = await litellm.acompletion(
-            model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4", 
-            messages=messages,
-            max_tokens=100,
-            temperature=0.7,
-            stream=True,
-        )
-        # Add any assertions here to check the response 
-        complete_response = "" 
-        async for chunk in response:
-            complete_response += chunk.choices[0].delta.content or "" 
-        print(f"complete_response: {complete_response}")
-        assert len(complete_response) > 0
-    except: 
-        pass
-
-asyncio.run(test_completion_chat_sagemaker())
-
 def test_completion_chat_sagemaker_mistral(): 
     try: 
         messages = [{"role": "user", "content": "Hey, how's it going?"}]
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 33199ea323f..23815784889 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -654,6 +654,28 @@ def test_sagemaker_weird_response():
         pytest.fail(f"An exception occurred - {str(e)}")
 # test_sagemaker_weird_response()
 
+@pytest.mark.asyncio
+async def test_sagemaker_streaming_async():
+    try: 
+        messages = [{"role": "user", "content": "Hey, how's it going?"}]
+        litellm.set_verbose=True
+        response = await litellm.acompletion(
+            model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4", 
+            messages=messages,
+            max_tokens=100,
+            temperature=0.7,
+            stream=True,
+        )
+
+        # Add any assertions here to check the response 
+        complete_response = "" 
+        async for chunk in response:
+            complete_response += chunk.choices[0].delta.content or "" 
+        print(f"complete_response: {complete_response}")
+        assert len(complete_response) > 0
+    except Exception as e: 
+        pytest.fail(f"An exception occurred - {str(e)}")
+
 # def test_completion_sagemaker_stream():
 #     try:
 #         response = completion(