diff --git a/airflow/providers/microsoft/azure/utils.py b/airflow/providers/microsoft/azure/utils.py index 1b738ed957a6a..a7a7e3896617a 100644 --- a/airflow/providers/microsoft/azure/utils.py +++ b/airflow/providers/microsoft/azure/utils.py @@ -59,8 +59,8 @@ def get_field(*, conn_id: str, conn_type: str, extras: dict, field_name: str): def _get_default_azure_credential( *, - managed_identity_client_id: str | None, - workload_identity_tenant_id: str | None, + managed_identity_client_id: str | None = None, + workload_identity_tenant_id: str | None = None, use_async: bool = False, ) -> DefaultAzureCredential | AsyncDefaultAzureCredential: """Get DefaultAzureCredential based on provided arguments. @@ -88,7 +88,7 @@ def _get_default_azure_credential( get_async_default_azure_credential: partial[AsyncDefaultAzureCredential] = partial( _get_default_azure_credential, # type: ignore[arg-type] - use_async=False, + use_async=True, ) diff --git a/tests/providers/microsoft/azure/test_utils.py b/tests/providers/microsoft/azure/test_utils.py index 5a081441cad35..f04acaab13e3c 100644 --- a/tests/providers/microsoft/azure/test_utils.py +++ b/tests/providers/microsoft/azure/test_utils.py @@ -25,7 +25,10 @@ from airflow.providers.microsoft.azure.utils import ( AzureIdentityCredentialAdapter, add_managed_identity_connection_widgets, + get_async_default_azure_credential, get_field, + # _get_default_azure_credential + get_sync_default_azure_credential, ) MODULE = "airflow.providers.microsoft.azure.utils" @@ -77,6 +80,19 @@ def test_func() -> dict[str, Any]: assert "workload_identity_tenant_id" in widgets +@mock.patch(f"{MODULE}.DefaultAzureCredential") +def test_get_sync_default_azure_credential(mock_default_azure_credential): + get_sync_default_azure_credential() + + assert mock_default_azure_credential.called + + +@mock.patch(f"{MODULE}.AsyncDefaultAzureCredential") +def test_get_async_default_azure_credential(mock_default_azure_credential): + get_async_default_azure_credential() + assert mock_default_azure_credential.called + + class TestAzureIdentityCredentialAdapter: @mock.patch(f"{MODULE}.PipelineRequest") @mock.patch(f"{MODULE}.BearerTokenCredentialPolicy")