Fix test case issue

litellm/llms/azure/common_utils.py

@@ -46,12 +46,11 @@ class AzureAuthResponse(BaseModel):
     Attributes:
         api_key: The API key for Azure OpenAI, if provided. Can be None.
         azure_ad_token_provider: A callable that provides an Azure AD token. Can be None.
-        azure_ad_token: The Azure AD token, if available. Can be None, a string token,
-                       or a callable that returns a token.
+        azure_ad_token: The Azure AD token, if available. Can be None.
     """
     api_key: Optional[str] = None
     azure_ad_token_provider: Optional[Callable[[], str]] = None
-    azure_ad_token: Union[None, str, Callable[[], str]] = None
+    azure_ad_token: Optional[str] = None
 
 
 def process_azure_headers(headers: Union[httpx.Headers, dict]) -> dict:
@@ -340,7 +339,7 @@ def get_azure_api_key_or_token(
             "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth"
         )
         try:
-            azure_ad_token = get_azure_ad_token_provider()
+            azure_ad_token_provider = get_azure_ad_token_provider()
         except ValueError:
             verbose_logger.debug("Azure AD Token Provider could not be used.")
 

tests/litellm/llms/azure/response/test_azure_responses_transformation.py

@@ -181,7 +181,7 @@ def test_validate_environment_with_service_principal(self, mock_azure_env_vars_n
         litellm_params = {}
 
         # Create a token provider that returns a token when called
-        mock_token = "service_principal_token"
+        mock_token = lambda : "service_principal_token"
 
         with patch("litellm.api_key", None):
             with patch("litellm.azure_key", None):
@@ -201,7 +201,7 @@ def test_validate_environment_with_service_principal(self, mock_azure_env_vars_n
 
                         # Verify the result contains the expected Authorization header
                         assert "Authorization" in result
-                        assert result["Authorization"] == f"Bearer {mock_token}"
+                        assert result["Authorization"] == f"Bearer {mock_token()}"
 
     def test_validate_environment_service_principal_error(self, mock_azure_env_vars_none):
         """Test validate_environment when Service Principal raises ValueError"""

tests/litellm/llms/azure/test_azure_common_utils.py

@@ -908,16 +908,6 @@ def test_AzureAuthResponse_with_correct_value():
     assert response.azure_ad_token_provider is None
     assert response.azure_ad_token == fake_token
 
-    azure_ad_token = lambda: fake_token
-    response = AzureAuthResponse(
-        azure_ad_token=azure_ad_token
-    )
-
-    assert response.api_key is None
-    assert response.azure_ad_token_provider is None
-    assert isinstance(response.azure_ad_token, Callable)
-    assert response.azure_ad_token() == fake_token
-
 def test_AzureAuthResponse_with_incorrect_value():
     with pytest.raises(ValidationError):
         response = AzureAuthResponse(