From 3af994a32a5ab3ea510448b59f5191460b679115 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Thu, 15 Jan 2026 19:21:31 +0530
Subject: [PATCH] Fix: test router

---
 tests/test_litellm/test_router.py | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 6279e96305f..ff9fe6b738a 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -1231,18 +1231,30 @@ async def test_acompletion_streaming_disable_fallbacks_midstream():
             return self
 
         async def __anext__(self):
-            if self.index >= len(self.items):
-                raise StopAsyncIteration
             if self.index == self.error_after_index:
                 raise self.error
+            if self.index >= len(self.items):
+                raise StopAsyncIteration
             item = self.items[self.index]
             self.index += 1
             self.chunks.append(item)
             return item
 
-    mock_chunks = [
-        MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]),
-    ]
+    # Create properly structured mock chunks using ModelResponse
+    from litellm.types.utils import Delta, ModelResponse, StreamingChoices
+
+    mock_chunk = ModelResponse(
+        id="chatcmpl-123",
+        choices=[
+            StreamingChoices(
+                index=0, delta=Delta(content="Hello", role="assistant"), finish_reason=None
+            )
+        ],
+        created=1234567890,
+        model="gpt-4",
+        object="chat.completion.chunk",
+    )
+    mock_chunks = [mock_chunk]
 
     mock_error_response = AsyncIteratorWithError(
         mock_chunks, 1, error_with_original