From a64d1de80f988fda56304db2cc74abac1a82d500 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 24 Jun 2024 21:13:01 -0700
Subject: [PATCH] test - aliases on /spend/calculate

---
 .../tests/test_spend_calculate_endpoint.py    | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/litellm/tests/test_spend_calculate_endpoint.py b/litellm/tests/test_spend_calculate_endpoint.py
index f8aff337ec..8bdd4a54d8 100644
--- a/litellm/tests/test_spend_calculate_endpoint.py
+++ b/litellm/tests/test_spend_calculate_endpoint.py
@@ -101,3 +101,41 @@ async def test_spend_calc_using_response():
     print("calculated cost", cost_obj)
     cost = cost_obj["cost"]
     assert cost > 0.0
+
+
+@pytest.mark.asyncio
+async def test_spend_calc_model_alias_on_router_messages():
+    from litellm.proxy.proxy_server import llm_router as init_llm_router
+
+    temp_llm_router = Router(
+        model_list=[
+            {
+                "model_name": "gpt-4o",
+                "litellm_params": {
+                    "model": "gpt-4o",
+                },
+            }
+        ],
+        model_group_alias={
+            "gpt4o": "gpt-4o",
+        },
+    )
+
+    setattr(litellm.proxy.proxy_server, "llm_router", temp_llm_router)
+
+    cost_obj = await calculate_spend(
+        request=SpendCalculateRequest(
+            model="gpt4o",
+            messages=[
+                {"role": "user", "content": "What is the capital of France?"},
+            ],
+        )
+    )
+
+    print("calculated cost", cost_obj)
+    _cost = cost_obj["cost"]
+
+    assert _cost > 0.0
+
+    # set router to init value
+    setattr(litellm.proxy.proxy_server, "llm_router", init_llm_router)