From 6fee7eff0786ede846f8d95fbc701617e052d92b Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Mon, 29 Sep 2025 13:15:43 -0700
Subject: [PATCH] test_litellm_overhead_stream

---
 .../test_litellm_overhead.py                  | 46 +++++++++----------
 1 file changed, 22 insertions(+), 24 deletions(-)

diff --git a/tests/litellm_utils_tests/test_litellm_overhead.py b/tests/litellm_utils_tests/test_litellm_overhead.py
index 6ab5eaf6174..6e46e935463 100644
--- a/tests/litellm_utils_tests/test_litellm_overhead.py
+++ b/tests/litellm_utils_tests/test_litellm_overhead.py
@@ -23,7 +23,7 @@ import litellm
         "vertex_ai/gemini-1.0-pro-vision-001",
     ],
 )
-async def test_litellm_overhead(model):
+async def test_litellm_overhead_non_streaming(model):
     """
     - Test we can see the litellm overhead and that it is less than 40% of the total request time
     """
@@ -37,20 +37,15 @@ async def test_litellm_overhead(model):
     #########################################################
     # Specific cases for models
     #########################################################
-    if model == "vertex_ai/gemini-1.0-pro-vision-001":
+    if model == "vertex_ai/gemini-1.0-pro-vision-001" or model == "openai/self_hosted":
         kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/"
         # warmup call for auth validation on vertex_ai models
         await litellm.acompletion(**kwargs)
-    if model == "openai/self_hosted":
-        response = await litellm.acompletion(
-            model=model,
-            messages=[{"role": "user", "content": "Hello, world!"}],
-            api_base="https://exampleopenaiendpoint-production.up.railway.app/",
-        )
-    else:
-        response = await litellm.acompletion(
-            **kwargs
-        )
+
+
+    response = await litellm.acompletion(
+        **kwargs
+    )
     #########################################################
     # End of specific cases for models
     #########################################################
@@ -92,19 +87,22 @@ async def test_litellm_overhead_stream(model):
 
     litellm._turn_on_debug()
     start_time = datetime.now()
+    kwargs ={
+        "messages": [{"role": "user", "content": "Hello, world!"}],
+        "model": model,
+        "stream": True,
+    }
+    #########################################################
+    # Specific cases for models
+    #########################################################
     if model == "openai/self_hosted":
-        response = await litellm.acompletion(
-            model=model,
-            messages=[{"role": "user", "content": "Hello, world!"}],
-            api_base="https://exampleopenaiendpoint-production.up.railway.app/",
-            stream=True,
-        )
-    else:
-        response = await litellm.acompletion(
-            model=model,
-            messages=[{"role": "user", "content": "Hello, world!"}],
-            stream=True,
-        )
+        kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/"
+        # warmup call for auth validation on vertex_ai models
+        await litellm.acompletion(**kwargs)
+    
+    response = await litellm.acompletion(
+        **kwargs
+    )
 
     async for chunk in response:
         print()