From 43de4ffeab7363adc35be368ad824e8c5b61e4fc Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Tue, 13 Jan 2026 13:50:07 -0800
Subject: [PATCH] fix - azure model router integration

---
 tests/llm_translation/test_azure_ai.py | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py
index 8d1ebd95904..2977c50a63f 100644
--- a/tests/llm_translation/test_azure_ai.py
+++ b/tests/llm_translation/test_azure_ai.py
@@ -364,4 +364,22 @@ def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base):
     except litellm.Timeout as e:
         pass
     except Exception as e:
-        pytest.fail(f"Error occurred: {e}")
\ No newline at end of file
+        pytest.fail(f"Error occurred: {e}")
+
+
+@pytest.mark.asyncio
+async def test_azure_ai_model_router():
+    litellm._turn_on_debug()
+    response = await litellm.acompletion(
+        model="azure_ai/azure-model-router",
+        messages=[{"role": "user", "content": "hi who is this"}],
+        api_base="https://ishaa-mh6uutut-swedencentral.cognitiveservices.azure.com/openai/v1/",
+        api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
+    )
+    print("response: ", response)
+    
+
+    # Check response cost
+    tracked_cost = response._hidden_params["response_cost"]
+    assert tracked_cost > 0
+    print("Tracked cost: ", tracked_cost)
\ No newline at end of file