From 658a23c3f15dc2e56577c19d6d6151f57fc67f76 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 11 Feb 2026 10:30:56 -0800
Subject: [PATCH] test_azure_ai_validate_environment_with_api_key

---
 .../chat/test_azure_ai_transformation.py      | 52 +++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py
index c91ef31bba5..d903d7c85f1 100644
--- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py
+++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py
@@ -45,6 +45,58 @@ def test_azure_ai_validate_environment():
     assert headers["Content-Type"] == "application/json"
 
 
+def test_azure_ai_validate_environment_with_api_key():
+    """
+    Test that when api_key is provided, it is set in the api-key header
+    for Azure Foundry endpoints (.services.ai.azure.com).
+    """
+    config = AzureAIStudioConfig()
+    headers = config.validate_environment(
+        headers={},
+        model="Kimi-K2.5",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key="test-api-key",
+        api_base="https://my-endpoint.services.ai.azure.com",
+    )
+    assert headers["api-key"] == "test-api-key"
+    assert headers["Content-Type"] == "application/json"
+
+
+def test_azure_ai_validate_environment_with_azure_ad_token():
+    """
+    Test that when no api_key is provided but Azure AD credentials are available,
+    the Authorization header is set with a Bearer token.
+
+    Regression test for https://github.com/BerriAI/litellm/issues/20759
+    """
+    import litellm
+
+    config = AzureAIStudioConfig()
+    with patch(
+        "litellm.llms.azure.common_utils.get_azure_ad_token",
+        return_value="fake-azure-ad-token",
+    ), patch(
+        "litellm.llms.azure.common_utils.get_secret_str",
+        return_value=None,
+    ), patch.object(litellm, "api_key", None), patch.object(
+        litellm, "azure_key", None
+    ):
+        headers = config.validate_environment(
+            headers={},
+            model="Kimi-K2.5",
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            api_key=None,
+            api_base="https://my-endpoint.services.ai.azure.com",
+        )
+    assert headers.get("Authorization") == "Bearer fake-azure-ad-token"
+    assert "api-key" not in headers
+    assert headers["Content-Type"] == "application/json"
+
+
 def test_azure_ai_grok_stop_parameter_handling():
     """
     Test that Grok models properly handle stop parameter filtering in Azure AI Studio.