From 6d63c0015b1050e98c40f6469e96214b3fa42018 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 27 Dec 2023 15:10:43 +0530
Subject: [PATCH] test(test_router_fallbacks.py): add testing for sync
 streaming fallbacks

---
 litellm/tests/test_router_fallbacks.py | 79 ++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/litellm/tests/test_router_fallbacks.py b/litellm/tests/test_router_fallbacks.py
index 6f97cabec..e601343e5 100644
--- a/litellm/tests/test_router_fallbacks.py
+++ b/litellm/tests/test_router_fallbacks.py
@@ -482,3 +482,82 @@ async def test_async_fallbacks_streaming():
         pytest.fail(f"An exception occurred: {e}")
     finally:
         router.reset()
+
+
+def test_sync_fallbacks_streaming():
+    try:
+        model_list = [
+            {  # list of model deployments
+                "model_name": "azure/gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "azure/chatgpt-v-2",
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                },
+                "tpm": 240000,
+                "rpm": 1800,
+            },
+            {  # list of model deployments
+                "model_name": "azure/gpt-3.5-turbo-context-fallback",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "azure/chatgpt-v-2",
+                    "api_key": os.getenv("AZURE_API_KEY"),
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                },
+                "tpm": 240000,
+                "rpm": 1800,
+            },
+            {
+                "model_name": "azure/gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "azure/chatgpt-functioncalling",
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                },
+                "tpm": 240000,
+                "rpm": 1800,
+            },
+            {
+                "model_name": "gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "gpt-3.5-turbo",
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000,
+            },
+            {
+                "model_name": "gpt-3.5-turbo-16k",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "gpt-3.5-turbo-16k",
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000,
+            },
+        ]
+
+        litellm.set_verbose = True
+        customHandler = MyCustomHandler()
+        litellm.callbacks = [customHandler]
+        router = Router(
+            model_list=model_list,
+            fallbacks=[{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}],
+            context_window_fallbacks=[
+                {"azure/gpt-3.5-turbo-context-fallback": ["gpt-3.5-turbo-16k"]},
+                {"gpt-3.5-turbo": ["gpt-3.5-turbo-16k"]},
+            ],
+            set_verbose=False,
+        )
+        response = router.completion(**kwargs, stream=True)
+        print(f"response: {response}")
+        time.sleep(0.05)  # allow a delay as success_callbacks are on a separate thread
+        assert customHandler.previous_models == 1  # 0 retries, 1 fallback
+
+        print("Passed ! Test router_fallbacks: test_sync_fallbacks()")
+        router.reset()
+    except Exception as e:
+        print(e)