From e5e477d7f54cd513d446cb2368588632bc79d2c9 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 7 May 2024 12:19:44 -0700
Subject: [PATCH] test - lowest cost routing

---
 litellm/tests/test_lowest_cost_routing.py | 70 +++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100644 litellm/tests/test_lowest_cost_routing.py

diff --git a/litellm/tests/test_lowest_cost_routing.py b/litellm/tests/test_lowest_cost_routing.py
new file mode 100644
index 00000000000..1fa6c3e63bb
--- /dev/null
+++ b/litellm/tests/test_lowest_cost_routing.py
@@ -0,0 +1,70 @@
+#### What this tests ####
+#    This tests the router's ability to pick deployment with lowest latency
+
+import sys, os, asyncio, time, random
+from datetime import datetime
+import traceback
+from dotenv import load_dotenv
+
+load_dotenv()
+import os, copy
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import pytest
+from litellm import Router
+from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
+from litellm.caching import DualCache
+
+### UNIT TESTS FOR LATENCY ROUTING ###
+
+
+def test_get_available_deployments():
+    test_cache = DualCache()
+    model_list = [
+        {
+            "model_name": "gpt-3.5-turbo",
+            "litellm_params": {"model": "gpt-4"},
+            "model_info": {"id": "openai-gpt-4"},
+        },
+        {
+            "model_name": "gpt-3.5-turbo",
+            "litellm_params": {"model": "groq/llama3-8b-8192"},
+            "model_info": {"id": "groq-llama"},
+        },
+    ]
+    lowest_cost_logger = LowestCostLoggingHandler(
+        router_cache=test_cache, model_list=model_list
+    )
+    model_group = "gpt-3.5-turbo"
+
+    ## CHECK WHAT'S SELECTED ##
+    selected_model = lowest_cost_logger.get_available_deployments(
+        model_group=model_group, healthy_deployments=model_list
+    )
+    print("selected model: ", selected_model)
+
+    assert selected_model["model_info"]["id"] == "groq-llama"
+
+
+async def _deploy(lowest_latency_logger, deployment_id, tokens_used, duration):
+    kwargs = {
+        "litellm_params": {
+            "metadata": {
+                "model_group": "gpt-3.5-turbo",
+                "deployment": "azure/chatgpt-v-2",
+            },
+            "model_info": {"id": deployment_id},
+        }
+    }
+    start_time = time.time()
+    response_obj = {"usage": {"total_tokens": tokens_used}}
+    time.sleep(duration)
+    end_time = time.time()
+    lowest_latency_logger.log_success_event(
+        response_obj=response_obj,
+        kwargs=kwargs,
+        start_time=start_time,
+        end_time=end_time,
+    )