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

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()