From 4007785b66c8bebf4528325a855c8aee96166b33 Mon Sep 17 00:00:00 2001 From: Giles Odigwe <79032838+giles17@users.noreply.github.com> Date: Thu, 14 Aug 2025 09:19:44 -0700 Subject: [PATCH] Azure package unit tests to improve coverage (#414) Co-authored-by: Giles Odigwe --- .../tests/test_azure_assistants_client.py | 183 +++++++++++++++++- .../tests/test_entra_id_authentication.py | 156 +++++++++++++++ 2 files changed, 338 insertions(+), 1 deletion(-) create mode 100644 python/packages/azure/tests/test_entra_id_authentication.py diff --git a/python/packages/azure/tests/test_azure_assistants_client.py b/python/packages/azure/tests/test_azure_assistants_client.py index 59bdae240e..45372664b0 100644 --- a/python/packages/azure/tests/test_azure_assistants_client.py +++ b/python/packages/azure/tests/test_azure_assistants_client.py @@ -2,7 +2,7 @@ import os from typing import Annotated -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from agent_framework import ( @@ -383,3 +383,184 @@ async def test_azure_assistants_client_with_existing_assistant() -> None: assert response is not None assert isinstance(response, ChatResponse) assert len(response.text) > 0 + + +def test_azure_assistants_client_entra_id_authentication() -> None: + """Test Entra ID authentication path with ad_credential.""" + mock_credential = MagicMock() + + with ( + patch("agent_framework_azure._assistants_client.AzureOpenAISettings") as mock_settings_class, + patch("agent_framework_azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client, + patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None), + ): + mock_settings = MagicMock() + mock_settings.chat_deployment_name = "test-deployment" + mock_settings.api_key = None # No API key to trigger Entra ID path + mock_settings.token_endpoint = "https://login.microsoftonline.com/test" + mock_settings.get_azure_auth_token.return_value = "entra-token-12345" + mock_settings.api_version = "2024-05-01-preview" + mock_settings.endpoint = "https://test-endpoint.openai.azure.com" + mock_settings.base_url = None + mock_settings_class.return_value = mock_settings + + client = AzureAssistantsClient( + deployment_name="test-deployment", + api_key="placeholder-key", + endpoint="https://test-endpoint.openai.azure.com", + ad_credential=mock_credential, + token_endpoint="https://login.microsoftonline.com/test", + ) + + # Verify Entra ID token was requested + mock_settings.get_azure_auth_token.assert_called_once_with(mock_credential) + + # Verify client was created with the token + mock_azure_client.assert_called_once() + call_args = mock_azure_client.call_args[1] + assert call_args["azure_ad_token"] == "entra-token-12345" + + assert client is not None + assert isinstance(client, AzureAssistantsClient) + + +def test_azure_assistants_client_no_authentication_error() -> None: + """Test authentication validation error when no auth provided.""" + with patch("agent_framework_azure._assistants_client.AzureOpenAISettings") as mock_settings_class: + mock_settings = MagicMock() + mock_settings.chat_deployment_name = "test-deployment" + mock_settings.api_key = None # No API key + mock_settings.token_endpoint = None # No token endpoint + mock_settings_class.return_value = mock_settings + + # Test missing authentication raises error + with pytest.raises(ServiceInitializationError, match="API key, ad_token, or ad_token_provider is required"): + AzureAssistantsClient( + deployment_name="test-deployment", + endpoint="https://test-endpoint.openai.azure.com", + # No authentication provided at all + ) + + +def test_azure_assistants_client_ad_token_authentication() -> None: + """Test ad_token authentication client parameter path.""" + with ( + patch("agent_framework_azure._assistants_client.AzureOpenAISettings") as mock_settings_class, + patch("agent_framework_azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client, + patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None), + ): + mock_settings = MagicMock() + mock_settings.chat_deployment_name = "test-deployment" + mock_settings.api_key = None # No API key + mock_settings.api_version = "2024-05-01-preview" + mock_settings.endpoint = "https://test-endpoint.openai.azure.com" + mock_settings.base_url = None + mock_settings_class.return_value = mock_settings + + client = AzureAssistantsClient( + deployment_name="test-deployment", + endpoint="https://test-endpoint.openai.azure.com", + ad_token="test-ad-token-12345", + ) + + # ad_token path + mock_azure_client.assert_called_once() + call_args = mock_azure_client.call_args[1] + assert call_args["azure_ad_token"] == "test-ad-token-12345" + + assert client is not None + assert isinstance(client, AzureAssistantsClient) + + +def test_azure_assistants_client_ad_token_provider_authentication() -> None: + """Test ad_token_provider authentication client parameter path.""" + from openai.lib.azure import AsyncAzureADTokenProvider + + mock_token_provider = MagicMock(spec=AsyncAzureADTokenProvider) + + with ( + patch("agent_framework_azure._assistants_client.AzureOpenAISettings") as mock_settings_class, + patch("agent_framework_azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client, + patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None), + ): + mock_settings = MagicMock() + mock_settings.chat_deployment_name = "test-deployment" + mock_settings.api_key = None # No API key + mock_settings.api_version = "2024-05-01-preview" + mock_settings.endpoint = "https://test-endpoint.openai.azure.com" + mock_settings.base_url = None + mock_settings_class.return_value = mock_settings + + client = AzureAssistantsClient( + deployment_name="test-deployment", + endpoint="https://test-endpoint.openai.azure.com", + ad_token_provider=mock_token_provider, + ) + + # ad_token_provider path + mock_azure_client.assert_called_once() + call_args = mock_azure_client.call_args[1] + assert call_args["azure_ad_token_provider"] is mock_token_provider + + assert client is not None + assert isinstance(client, AzureAssistantsClient) + + +def test_azure_assistants_client_base_url_configuration() -> None: + """Test base_url client parameter path.""" + with ( + patch("agent_framework_azure._assistants_client.AzureOpenAISettings") as mock_settings_class, + patch("agent_framework_azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client, + patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None), + ): + mock_settings = MagicMock() + mock_settings.chat_deployment_name = "test-deployment" + mock_settings.api_key.get_secret_value.return_value = "test-api-key" + mock_settings.base_url = "https://custom-base-url.com" + mock_settings.endpoint = None # No endpoint, should use base_url + mock_settings.api_version = "2024-05-01-preview" + mock_settings_class.return_value = mock_settings + + client = AzureAssistantsClient( + deployment_name="test-deployment", api_key="test-api-key", base_url="https://custom-base-url.com" + ) + + # base_url path + mock_azure_client.assert_called_once() + call_args = mock_azure_client.call_args[1] + assert call_args["base_url"] == "https://custom-base-url.com" + assert "azure_endpoint" not in call_args + + assert client is not None + assert isinstance(client, AzureAssistantsClient) + + +def test_azure_assistants_client_azure_endpoint_configuration() -> None: + """Test azure_endpoint client parameter path.""" + with ( + patch("agent_framework_azure._assistants_client.AzureOpenAISettings") as mock_settings_class, + patch("agent_framework_azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client, + patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None), + ): + mock_settings = MagicMock() + mock_settings.chat_deployment_name = "test-deployment" + mock_settings.api_key.get_secret_value.return_value = "test-api-key" + mock_settings.base_url = None # No base_url + mock_settings.endpoint = "https://test-endpoint.openai.azure.com" + mock_settings.api_version = "2024-05-01-preview" + mock_settings_class.return_value = mock_settings + + client = AzureAssistantsClient( + deployment_name="test-deployment", + api_key="test-api-key", + endpoint="https://test-endpoint.openai.azure.com", + ) + + # azure_endpoint path + mock_azure_client.assert_called_once() + call_args = mock_azure_client.call_args[1] + assert call_args["azure_endpoint"] == "https://test-endpoint.openai.azure.com" + assert "base_url" not in call_args + + assert client is not None + assert isinstance(client, AzureAssistantsClient) diff --git a/python/packages/azure/tests/test_entra_id_authentication.py b/python/packages/azure/tests/test_entra_id_authentication.py new file mode 100644 index 0000000000..5b894fce06 --- /dev/null +++ b/python/packages/azure/tests/test_entra_id_authentication.py @@ -0,0 +1,156 @@ +# Copyright (c) Microsoft. All rights reserved. + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from agent_framework.exceptions import ServiceInvalidAuthError +from azure.core.exceptions import ClientAuthenticationError + +from agent_framework_azure._entra_id_authentication import ( + get_entra_auth_token, + get_entra_auth_token_async, +) + + +@pytest.fixture +def mock_credential() -> MagicMock: + """Mock synchronous ChainedTokenCredential.""" + mock_cred = MagicMock() + # Create a mock token object with a .token attribute + mock_token = MagicMock() + mock_token.token = "test-access-token-12345" + mock_cred.get_token.return_value = mock_token + return mock_cred + + +@pytest.fixture +def mock_async_credential() -> MagicMock: + """Mock asynchronous ChainedTokenCredential.""" + mock_cred = MagicMock() + # Create a mock token object with a .token attribute + mock_token = MagicMock() + mock_token.token = "test-async-access-token-12345" + mock_cred.get_token = AsyncMock(return_value=mock_token) + return mock_cred + + +def test_get_entra_auth_token_success(mock_credential: MagicMock) -> None: + """Test successful token retrieval with sync function.""" + + token_endpoint = "https://test-endpoint.com/.default" + + result = get_entra_auth_token(mock_credential, token_endpoint) + + # Assert - check the results + assert result == "test-access-token-12345" + mock_credential.get_token.assert_called_once_with(token_endpoint) + + +async def test_get_entra_auth_token_async_success(mock_async_credential: MagicMock) -> None: + """Test successful token retrieval with async function.""" + + token_endpoint = "https://test-endpoint.com/.default" + + result = await get_entra_auth_token_async(mock_async_credential, token_endpoint) + + # Assert - check the results + assert result == "test-async-access-token-12345" + mock_async_credential.get_token.assert_called_once_with(token_endpoint) + + +def test_get_entra_auth_token_missing_endpoint(mock_credential: MagicMock) -> None: + """Test that missing token endpoint raises ServiceInvalidAuthError.""" + # Test with empty string + with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"): + get_entra_auth_token(mock_credential, "") + + # Test with None + with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"): + get_entra_auth_token(mock_credential, None) # type: ignore + + +async def test_get_entra_auth_token_async_missing_endpoint(mock_async_credential: MagicMock) -> None: + """Test that missing token endpoint raises ServiceInvalidAuthError in async function.""" + # Test with empty string + with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"): + await get_entra_auth_token_async(mock_async_credential, "") + + # Test with None + with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"): + await get_entra_auth_token_async(mock_async_credential, None) # type: ignore + + +def test_get_entra_auth_token_auth_failure(mock_credential: MagicMock) -> None: + """Test that Azure authentication failure returns None.""" + + mock_credential.get_token.side_effect = ClientAuthenticationError("Auth failed") + token_endpoint = "https://test-endpoint.com/.default" + + result = get_entra_auth_token(mock_credential, token_endpoint) + + # Assert - should return None on auth failure + assert result is None + mock_credential.get_token.assert_called_once_with(token_endpoint) + + +async def test_get_entra_auth_token_async_auth_failure(mock_async_credential: MagicMock) -> None: + """Test that Azure authentication failure returns None in async function.""" + + mock_async_credential.get_token.side_effect = ClientAuthenticationError("Auth failed") + token_endpoint = "https://test-endpoint.com/.default" + + result = await get_entra_auth_token_async(mock_async_credential, token_endpoint) + + # Assert - should return None on auth failure + assert result is None + mock_async_credential.get_token.assert_called_once_with(token_endpoint) + + +def test_get_entra_auth_token_none_token_response(mock_credential: MagicMock) -> None: + """Test that None token response returns None.""" + mock_credential.get_token.return_value = None + token_endpoint = "https://test-endpoint.com/.default" + + result = get_entra_auth_token(mock_credential, token_endpoint) + + # Assert + assert result is None + mock_credential.get_token.assert_called_once_with(token_endpoint) + + +async def test_get_entra_auth_token_async_none_token_response(mock_async_credential: MagicMock) -> None: + """Test that None token response returns None in async function.""" + mock_async_credential.get_token.return_value = None + token_endpoint = "https://test-endpoint.com/.default" + + result = await get_entra_auth_token_async(mock_async_credential, token_endpoint) + + # Assert + assert result is None + mock_async_credential.get_token.assert_called_once_with(token_endpoint) + + +def test_get_entra_auth_token_with_kwargs(mock_credential: MagicMock) -> None: + """Test that kwargs are passed through to get_token.""" + + token_endpoint = "https://test-endpoint.com/.default" + extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"} + + result = get_entra_auth_token(mock_credential, token_endpoint, **extra_kwargs) + + # Assert + assert result == "test-access-token-12345" + mock_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs) + + +async def test_get_entra_auth_token_async_with_kwargs(mock_async_credential: MagicMock) -> None: + """Test that kwargs are passed through to async get_token.""" + + token_endpoint = "https://test-endpoint.com/.default" + extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"} + + result = await get_entra_auth_token_async(mock_async_credential, token_endpoint, **extra_kwargs) + + # Assert + assert result == "test-async-access-token-12345" + mock_async_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)