test(test_azure_common_utils.py): add unit testing for common azure client params function

litellm/llms/azure/azure.py

@@ -33,6 +33,7 @@ from .common_utils import (
     AzureOpenAIError,
     get_azure_ad_token_from_oidc,
     process_azure_headers,
+    select_azure_base_url_or_endpoint,
 )
 
 
@@ -98,18 +99,6 @@ class AzureOpenAIAssistantsAPIConfig:
         return optional_params
 
 
-def select_azure_base_url_or_endpoint(azure_client_params: dict):
-    azure_endpoint = azure_client_params.get("azure_endpoint", None)
-    if azure_endpoint is not None:
-        # see : https://github.com/openai/openai-python/blob/3d61ed42aba652b547029095a7eb269ad4e1e957/src/openai/lib/azure.py#L192
-        if "/openai/deployments" in azure_endpoint:
-            # this is base_url, not an azure_endpoint
-            azure_client_params["base_url"] = azure_endpoint
-            azure_client_params.pop("azure_endpoint")
-
-    return azure_client_params
-
-
 def _check_dynamic_azure_params(
     azure_client_params: dict,
     azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]],

litellm/llms/azure/common_utils.py

@@ -12,7 +12,6 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.secret_managers.get_azure_ad_token_provider import (
     get_azure_ad_token_provider,
 )
-from litellm.secret_managers.get_secret import get_secret
 from litellm.secret_managers.main import get_secret_str
 
 azure_ad_cache = DualCache()
@@ -202,7 +201,7 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str):
             message="AZURE_CLIENT_ID and AZURE_TENANT_ID must be set",
         )
 
-    oidc_token = get_secret(azure_ad_token)
+    oidc_token = get_secret_str(azure_ad_token)
 
     if oidc_token is None:
         raise AzureOpenAIError(
@@ -264,6 +263,18 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str):
     return azure_ad_token_access_token
 
 
+def select_azure_base_url_or_endpoint(azure_client_params: dict):
+    azure_endpoint = azure_client_params.get("azure_endpoint", None)
+    if azure_endpoint is not None:
+        # see : https://github.com/openai/openai-python/blob/3d61ed42aba652b547029095a7eb269ad4e1e957/src/openai/lib/azure.py#L192
+        if "/openai/deployments" in azure_endpoint:
+            # this is base_url, not an azure_endpoint
+            azure_client_params["base_url"] = azure_endpoint
+            azure_client_params.pop("azure_endpoint")
+
+    return azure_client_params
+
+
 def initialize_azure_sdk_client(
     litellm_params: dict,
     api_key: Optional[str],
@@ -324,8 +335,6 @@ def initialize_azure_sdk_client(
 
     if azure_ad_token_provider is not None:
         azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider
-    from litellm.llms.azure.azure import select_azure_base_url_or_endpoint
-
     # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client
     # required to support GPT-4 vision enhancements, since base_url needs to be set on Azure OpenAI Client
     azure_client_params = select_azure_base_url_or_endpoint(azure_client_params)

tests/litellm/llms/azure/test_azure_common_utils.py

@@ -0,0 +1,212 @@
+import json
+import os
+import sys
+from typing import Callable, Optional
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../../..")
+)  # Adds the parent directory to the system path
+import litellm
+from litellm.llms.azure.common_utils import initialize_azure_sdk_client
+
+
+# Mock the necessary dependencies
+@pytest.fixture
+def setup_mocks():
+    with patch(
+        "litellm.llms.azure.common_utils.get_azure_ad_token_from_entrata_id"
+    ) as mock_entrata_token, patch(
+        "litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password"
+    ) as mock_username_password_token, patch(
+        "litellm.llms.azure.common_utils.get_azure_ad_token_from_oidc"
+    ) as mock_oidc_token, patch(
+        "litellm.llms.azure.common_utils.get_azure_ad_token_provider"
+    ) as mock_token_provider, patch(
+        "litellm.llms.azure.common_utils.litellm"
+    ) as mock_litellm, patch(
+        "litellm.llms.azure.common_utils.verbose_logger"
+    ) as mock_logger, patch(
+        "litellm.llms.azure.common_utils.select_azure_base_url_or_endpoint"
+    ) as mock_select_url:
+
+        # Configure mocks
+        mock_litellm.AZURE_DEFAULT_API_VERSION = "2023-05-15"
+        mock_litellm.enable_azure_ad_token_refresh = False
+
+        mock_entrata_token.return_value = lambda: "mock-entrata-token"
+        mock_username_password_token.return_value = (
+            lambda: "mock-username-password-token"
+        )
+        mock_oidc_token.return_value = "mock-oidc-token"
+        mock_token_provider.return_value = lambda: "mock-default-token"
+
+        mock_select_url.side_effect = lambda params: params
+
+        yield {
+            "entrata_token": mock_entrata_token,
+            "username_password_token": mock_username_password_token,
+            "oidc_token": mock_oidc_token,
+            "token_provider": mock_token_provider,
+            "litellm": mock_litellm,
+            "logger": mock_logger,
+            "select_url": mock_select_url,
+        }
+
+
+def test_initialize_with_api_key(setup_mocks):
+    # Test with api_key provided
+    result = initialize_azure_sdk_client(
+        litellm_params={},
+        api_key="test-api-key",
+        api_base="https://test.openai.azure.com",
+        model_name="gpt-4",
+        api_version="2023-06-01",
+    )
+
+    # Verify expected result
+    assert result["api_key"] == "test-api-key"
+    assert result["azure_endpoint"] == "https://test.openai.azure.com"
+    assert result["api_version"] == "2023-06-01"
+    assert "azure_ad_token" in result
+    assert result["azure_ad_token"] is None
+
+
+def test_initialize_with_tenant_credentials(setup_mocks):
+    # Test with tenant_id, client_id, and client_secret provided
+    result = initialize_azure_sdk_client(
+        litellm_params={
+            "tenant_id": "test-tenant-id",
+            "client_id": "test-client-id",
+            "client_secret": "test-client-secret",
+        },
+        api_key=None,
+        api_base="https://test.openai.azure.com",
+        model_name="gpt-4",
+        api_version=None,
+    )
+
+    # Verify that get_azure_ad_token_from_entrata_id was called
+    setup_mocks["entrata_token"].assert_called_once_with(
+        tenant_id="test-tenant-id",
+        client_id="test-client-id",
+        client_secret="test-client-secret",
+    )
+
+    # Verify expected result
+    assert result["api_key"] is None
+    assert result["azure_endpoint"] == "https://test.openai.azure.com"
+    assert "azure_ad_token_provider" in result
+
+
+def test_initialize_with_username_password(setup_mocks):
+    # Test with azure_username, azure_password, and client_id provided
+    result = initialize_azure_sdk_client(
+        litellm_params={
+            "azure_username": "test-username",
+            "azure_password": "test-password",
+            "client_id": "test-client-id",
+        },
+        api_key=None,
+        api_base="https://test.openai.azure.com",
+        model_name="gpt-4",
+        api_version=None,
+    )
+
+    # Verify that get_azure_ad_token_from_username_password was called
+    setup_mocks["username_password_token"].assert_called_once_with(
+        azure_username="test-username",
+        azure_password="test-password",
+        client_id="test-client-id",
+    )
+
+    # Verify expected result
+    assert "azure_ad_token_provider" in result
+
+
+def test_initialize_with_oidc_token(setup_mocks):
+    # Test with azure_ad_token that starts with "oidc/"
+    result = initialize_azure_sdk_client(
+        litellm_params={"azure_ad_token": "oidc/test-token"},
+        api_key=None,
+        api_base="https://test.openai.azure.com",
+        model_name="gpt-4",
+        api_version=None,
+    )
+
+    # Verify that get_azure_ad_token_from_oidc was called
+    setup_mocks["oidc_token"].assert_called_once_with("oidc/test-token")
+
+    # Verify expected result
+    assert result["azure_ad_token"] == "mock-oidc-token"
+
+
+def test_initialize_with_enable_token_refresh(setup_mocks):
+    # Enable token refresh
+    setup_mocks["litellm"].enable_azure_ad_token_refresh = True
+
+    # Test with token refresh enabled
+    result = initialize_azure_sdk_client(
+        litellm_params={},
+        api_key=None,
+        api_base="https://test.openai.azure.com",
+        model_name="gpt-4",
+        api_version=None,
+    )
+
+    # Verify that get_azure_ad_token_provider was called
+    setup_mocks["token_provider"].assert_called_once()
+
+    # Verify expected result
+    assert "azure_ad_token_provider" in result
+
+
+def test_initialize_with_token_refresh_error(setup_mocks):
+    # Enable token refresh but make it raise an error
+    setup_mocks["litellm"].enable_azure_ad_token_refresh = True
+    setup_mocks["token_provider"].side_effect = ValueError("Token provider error")
+
+    # Test with token refresh enabled but raising error
+    result = initialize_azure_sdk_client(
+        litellm_params={},
+        api_key=None,
+        api_base="https://test.openai.azure.com",
+        model_name="gpt-4",
+        api_version=None,
+    )
+
+    # Verify error was logged
+    setup_mocks["logger"].debug.assert_any_call(
+        "Azure AD Token Provider could not be used."
+    )
+
+
+def test_api_version_from_env_var(setup_mocks):
+    # Test api_version from environment variable
+    with patch.dict(os.environ, {"AZURE_API_VERSION": "2023-07-01"}):
+        result = initialize_azure_sdk_client(
+            litellm_params={},
+            api_key="test-api-key",
+            api_base="https://test.openai.azure.com",
+            model_name="gpt-4",
+            api_version=None,
+        )
+
+    # Verify expected result
+    assert result["api_version"] == "2023-07-01"
+
+
+def test_select_azure_base_url_called(setup_mocks):
+    # Test that select_azure_base_url_or_endpoint is called
+    result = initialize_azure_sdk_client(
+        litellm_params={},
+        api_key="test-api-key",
+        api_base="https://test.openai.azure.com",
+        model_name="gpt-4",
+        api_version="2023-06-01",
+    )
+
+    # Verify that select_azure_base_url_or_endpoint was called
+    setup_mocks["select_url"].assert_called_once()