mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Unify Azure credential handling across all Python packages (#4088)
Replace ad_token, ad_token_provider, and get_entra_auth_token with a unified credential parameter across all Azure-related packages. Core changes: - Add AzureCredentialTypes (TokenCredential | AsyncTokenCredential) and AzureTokenProvider (Callable[[], str | Awaitable[str]]) type aliases - Add resolve_credential_to_token_provider() using azure.identity's get_bearer_token_provider for automatic token caching/refresh - Update AzureOpenAIChatClient, AzureOpenAIResponsesClient, and AzureOpenAIAssistantsClient to accept credential: AzureCredentialTypes | AzureTokenProvider - Remove ad_token, ad_token_provider params and get_entra_auth_token helpers Package updates: - azure-ai: Accept AzureCredentialTypes on AzureAIClient, AzureAIAgentClient, AzureAIProjectAgentProvider, AzureAIAgentsProvider - azure-ai-search: Accept AzureCredentialTypes on AzureAISearchContextProvider - purview: Accept AzureCredentialTypes | AzureTokenProvider on PurviewClient, PurviewPolicyMiddleware, PurviewChatPolicyMiddleware Fixes #3449 Fixes #3500 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
4c8f595019
commit
fd4e6e816c
@@ -551,12 +551,16 @@ async def test_azure_assistants_client_agent_level_tool_persistence():
|
||||
|
||||
|
||||
def test_azure_assistants_client_entra_id_authentication() -> None:
|
||||
"""Test Entra ID authentication path with credential."""
|
||||
"""Test credential authentication path with sync credential."""
|
||||
mock_credential = MagicMock()
|
||||
mock_provider = MagicMock(return_value="token-string")
|
||||
|
||||
with (
|
||||
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
|
||||
patch("agent_framework.azure._assistants_client.get_entra_auth_token") as mock_get_token,
|
||||
patch(
|
||||
"agent_framework.azure._assistants_client.resolve_credential_to_token_provider",
|
||||
return_value=mock_provider,
|
||||
) as mock_resolve,
|
||||
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
|
||||
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
|
||||
):
|
||||
@@ -564,28 +568,26 @@ def test_azure_assistants_client_entra_id_authentication() -> None:
|
||||
"chat_deployment_name": "test-deployment",
|
||||
"responses_deployment_name": None,
|
||||
"api_key": None,
|
||||
"token_endpoint": "https://login.microsoftonline.com/test",
|
||||
"token_endpoint": "https://cognitiveservices.azure.com/.default",
|
||||
"api_version": "2024-05-01-preview",
|
||||
"endpoint": "https://test-endpoint.openai.azure.com",
|
||||
"base_url": None,
|
||||
}
|
||||
mock_get_token.return_value = "entra-token-12345"
|
||||
|
||||
client = AzureOpenAIAssistantsClient(
|
||||
deployment_name="test-deployment",
|
||||
api_key="placeholder-key",
|
||||
endpoint="https://test-endpoint.openai.azure.com",
|
||||
credential=mock_credential,
|
||||
token_endpoint="https://login.microsoftonline.com/test",
|
||||
token_endpoint="https://cognitiveservices.azure.com/.default",
|
||||
)
|
||||
|
||||
# Verify Entra ID token was requested
|
||||
mock_get_token.assert_called_once_with(mock_credential, "https://login.microsoftonline.com/test")
|
||||
# Verify credential was resolved to a token provider
|
||||
mock_resolve.assert_called_once_with(mock_credential, "https://cognitiveservices.azure.com/.default")
|
||||
|
||||
# Verify client was created with the token
|
||||
# Verify client was created with the token provider
|
||||
mock_azure_client.assert_called_once()
|
||||
call_args = mock_azure_client.call_args[1]
|
||||
assert call_args["azure_ad_token"] == "entra-token-12345"
|
||||
assert call_args["azure_ad_token_provider"] is mock_provider
|
||||
|
||||
assert client is not None
|
||||
assert isinstance(client, AzureOpenAIAssistantsClient)
|
||||
@@ -605,7 +607,7 @@ def test_azure_assistants_client_no_authentication_error() -> None:
|
||||
}
|
||||
|
||||
# Test missing authentication raises error
|
||||
with pytest.raises(ServiceInitializationError, match="API key, ad_token, or ad_token_provider is required"):
|
||||
with pytest.raises(ServiceInitializationError, match="api_key, credential, or a client"):
|
||||
AzureOpenAIAssistantsClient(
|
||||
deployment_name="test-deployment",
|
||||
endpoint="https://test-endpoint.openai.azure.com",
|
||||
@@ -613,10 +615,16 @@ def test_azure_assistants_client_no_authentication_error() -> None:
|
||||
)
|
||||
|
||||
|
||||
def test_azure_assistants_client_ad_token_authentication() -> None:
|
||||
"""Test ad_token authentication client parameter path."""
|
||||
def test_azure_assistants_client_callable_credential() -> None:
|
||||
"""Test callable token provider as credential."""
|
||||
mock_provider = MagicMock(return_value="my-token")
|
||||
|
||||
with (
|
||||
patch("agent_framework.azure._assistants_client.load_settings") as mock_load_settings,
|
||||
patch(
|
||||
"agent_framework.azure._assistants_client.resolve_credential_to_token_provider",
|
||||
return_value=mock_provider,
|
||||
),
|
||||
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
|
||||
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
|
||||
):
|
||||
@@ -624,7 +632,7 @@ def test_azure_assistants_client_ad_token_authentication() -> None:
|
||||
"chat_deployment_name": "test-deployment",
|
||||
"responses_deployment_name": None,
|
||||
"api_key": None,
|
||||
"token_endpoint": None,
|
||||
"token_endpoint": "https://cognitiveservices.azure.com/.default",
|
||||
"api_version": "2024-05-01-preview",
|
||||
"endpoint": "https://test-endpoint.openai.azure.com",
|
||||
"base_url": None,
|
||||
@@ -633,49 +641,14 @@ def test_azure_assistants_client_ad_token_authentication() -> None:
|
||||
client = AzureOpenAIAssistantsClient(
|
||||
deployment_name="test-deployment",
|
||||
endpoint="https://test-endpoint.openai.azure.com",
|
||||
ad_token="test-ad-token-12345",
|
||||
credential=mock_provider,
|
||||
token_endpoint="https://cognitiveservices.azure.com/.default",
|
||||
)
|
||||
|
||||
# ad_token path
|
||||
# Verify client was created with the token provider
|
||||
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, AzureOpenAIAssistantsClient)
|
||||
|
||||
|
||||
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.load_settings") as mock_load_settings,
|
||||
patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
|
||||
patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
|
||||
):
|
||||
mock_load_settings.return_value = {
|
||||
"chat_deployment_name": "test-deployment",
|
||||
"responses_deployment_name": None,
|
||||
"api_key": None,
|
||||
"token_endpoint": None,
|
||||
"api_version": "2024-05-01-preview",
|
||||
"endpoint": "https://test-endpoint.openai.azure.com",
|
||||
"base_url": None,
|
||||
}
|
||||
|
||||
client = AzureOpenAIAssistantsClient(
|
||||
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 call_args["azure_ad_token_provider"] is mock_provider
|
||||
|
||||
assert client is not None
|
||||
assert isinstance(client, AzureOpenAIAssistantsClient)
|
||||
|
||||
@@ -1,156 +1,61 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from azure.core.exceptions import ClientAuthenticationError
|
||||
from azure.core.credentials import TokenCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
from agent_framework.azure._entra_id_authentication import (
|
||||
get_entra_auth_token,
|
||||
get_entra_auth_token_async,
|
||||
resolve_credential_to_token_provider,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInvalidAuthError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_credential() -> MagicMock:
|
||||
"""Mock synchronous TokenCredential."""
|
||||
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
|
||||
TOKEN_ENDPOINT = "https://cognitiveservices.azure.com/.default"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_async_credential() -> MagicMock:
|
||||
"""Mock asynchronous AsyncTokenCredential."""
|
||||
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_resolve_sync_credential_returns_provider() -> None:
|
||||
"""Test that a sync TokenCredential is resolved via azure.identity.get_bearer_token_provider."""
|
||||
mock_credential = MagicMock(spec=TokenCredential)
|
||||
mock_provider = MagicMock(return_value="token-string")
|
||||
|
||||
with patch("azure.identity.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
|
||||
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
|
||||
|
||||
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
|
||||
assert result is mock_provider
|
||||
|
||||
|
||||
def test_get_entra_auth_token_success(mock_credential: MagicMock) -> None:
|
||||
"""Test successful token retrieval with sync function."""
|
||||
def test_resolve_async_credential_returns_provider() -> None:
|
||||
"""Test that an AsyncTokenCredential is resolved via azure.identity.aio.get_bearer_token_provider."""
|
||||
mock_credential = MagicMock(spec=AsyncTokenCredential)
|
||||
mock_provider = MagicMock(return_value="token-string")
|
||||
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
with patch("azure.identity.aio.get_bearer_token_provider", return_value=mock_provider) as mock_gbtp:
|
||||
result = resolve_credential_to_token_provider(mock_credential, TOKEN_ENDPOINT)
|
||||
|
||||
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)
|
||||
mock_gbtp.assert_called_once_with(mock_credential, TOKEN_ENDPOINT)
|
||||
assert result is mock_provider
|
||||
|
||||
|
||||
async def test_get_entra_auth_token_async_success(mock_async_credential: MagicMock) -> None:
|
||||
"""Test successful token retrieval with async function."""
|
||||
def test_resolve_callable_provider_passthrough() -> None:
|
||||
"""Test that a callable token provider is returned as-is, without needing token_endpoint."""
|
||||
my_provider = lambda: "my-token" # noqa: E731
|
||||
|
||||
token_endpoint = "https://test-endpoint.com/.default"
|
||||
# Works with token_endpoint
|
||||
assert resolve_credential_to_token_provider(my_provider, TOKEN_ENDPOINT) is my_provider
|
||||
|
||||
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)
|
||||
# Also works without token_endpoint
|
||||
assert resolve_credential_to_token_provider(my_provider, None) is my_provider
|
||||
assert resolve_credential_to_token_provider(my_provider, "") is my_provider
|
||||
|
||||
|
||||
def test_get_entra_auth_token_missing_endpoint(mock_credential: MagicMock) -> None:
|
||||
def test_resolve_missing_endpoint_raises() -> None:
|
||||
"""Test that missing token endpoint raises ServiceInvalidAuthError."""
|
||||
# Test with empty string
|
||||
mock_credential = MagicMock(spec=TokenCredential)
|
||||
|
||||
with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
|
||||
get_entra_auth_token(mock_credential, "")
|
||||
resolve_credential_to_token_provider(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)
|
||||
resolve_credential_to_token_provider(mock_credential, None) # type: ignore[arg-type]
|
||||
|
||||
Reference in New Issue
Block a user