From dc4235e7f1dd2d40c879cff52c8616bea4cf7efc Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 14 Nov 2024 18:56:59 +0530
Subject: [PATCH] test: update tests

---
 .../test_stream_chunk_builder.py              | 31 ++++++++++++-------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py
index 5fbdf07b8..a65d440a7 100644
--- a/tests/local_testing/test_stream_chunk_builder.py
+++ b/tests/local_testing/test_stream_chunk_builder.py
@@ -172,6 +172,8 @@ def test_stream_chunk_builder_litellm_usage_chunks():
     """
     Checks if stream_chunk_builder is able to correctly rebuild with given metadata from streaming chunks
     """
+    from litellm.types.utils import Usage
+
     messages = [
         {"role": "user", "content": "Tell me the funniest joke you know."},
         {
@@ -183,23 +185,28 @@ def test_stream_chunk_builder_litellm_usage_chunks():
         {"role": "user", "content": "\nI am waiting...\n\n...\n"},
     ]
     # make a regular gemini call
-    response = completion(
-        model="gemini/gemini-1.5-flash",
-        messages=messages,
-    )
 
-    usage: litellm.Usage = response.usage
+    usage: litellm.Usage = Usage(
+        completion_tokens=64,
+        prompt_tokens=55,
+        total_tokens=119,
+        completion_tokens_details=None,
+        prompt_tokens_details=None,
+    )
 
     gemini_pt = usage.prompt_tokens
 
     # make a streaming gemini call
-    response = completion(
-        model="gemini/gemini-1.5-flash",
-        messages=messages,
-        stream=True,
-        complete_response=True,
-        stream_options={"include_usage": True},
-    )
+    try:
+        response = completion(
+            model="gemini/gemini-1.5-flash",
+            messages=messages,
+            stream=True,
+            complete_response=True,
+            stream_options={"include_usage": True},
+        )
+    except litellm.ServiceUnavailableError as e:
+        pytest.skip(f"ServiceUnavailableError - {str(e)}")
 
     usage: litellm.Usage = response.usage