tests

sdk/identity/azure-identity/tests/test_chained_credential.py

@@ -3,9 +3,9 @@
 # Licensed under the MIT License.
 # ------------------------------------
 try:
-    from unittest.mock import Mock
+    from unittest.mock import MagicMock, Mock
 except ImportError:  # python < 3.3
-    from mock import Mock  # type: ignore
+    from mock import MagicMock, Mock  # type: ignore
 
 from azure.core.credentials import AccessToken
 from azure.core.exceptions import ClientAuthenticationError
@@ -17,6 +17,37 @@
 import pytest
 
 
+def test_close():
+    credentials = [MagicMock(close=Mock()) for _ in range(5)]
+    chain = ChainedTokenCredential(*credentials)
+
+    for credential in credentials:
+        assert credential.__exit__.call_count == 0
+
+    chain.close()
+
+    for credential in credentials:
+        assert credential.__exit__.call_count == 1
+
+
+def test_context_manager():
+    credentials = [MagicMock() for _ in range(5)]
+    chain = ChainedTokenCredential(*credentials)
+
+    for credential in credentials:
+        assert credential.__enter__.call_count == 0
+        assert credential.__exit__.call_count == 0
+
+    with chain:
+        for credential in credentials:
+            assert credential.__enter__.call_count == 1
+            assert credential.__exit__.call_count == 0
+
+    for credential in credentials:
+        assert credential.__enter__.call_count == 1
+        assert credential.__exit__.call_count == 1
+
+
 def test_error_message():
     first_error = "first_error"
     first_credential = Mock(

sdk/identity/azure-identity/tests/test_context_manager.py

@@ -0,0 +1,106 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+try:
+    from unittest.mock import MagicMock, patch
+except ImportError:
+    from mock import MagicMock, patch  # type: ignore
+
+from azure.identity import (
+    AzureApplicationCredential,
+    AzureCliCredential,
+    AzurePowerShellCredential,
+    AuthorizationCodeCredential,
+    CertificateCredential,
+    ClientSecretCredential,
+    DeviceCodeCredential,
+    EnvironmentCredential,
+    InteractiveBrowserCredential,
+    SharedTokenCacheCredential,
+    UsernamePasswordCredential,
+    VisualStudioCodeCredential,
+)
+from azure.identity._constants import EnvironmentVariables
+
+import pytest
+
+from test_certificate_credential import CERT_PATH
+from test_vscode_credential import GET_USER_SETTINGS
+
+
+class CredentialFixture:
+    def __init__(self, cls, default_kwargs=None, ctor_patch=None, get_token_patch=None):
+        self.cls = cls
+        self.get_token_patch = get_token_patch or MagicMock()
+        self._default_kwargs = default_kwargs or {}
+        self._ctor_patch = ctor_patch or MagicMock()
+
+    def __call__(self, **kwargs):
+        with self._ctor_patch:
+            return self.cls(**dict(self._default_kwargs, **kwargs))
+
+
+FIXTURES = (
+    CredentialFixture(
+        AuthorizationCodeCredential,
+        {kwarg: "..." for kwarg in ("tenant_id", "client_id", "authorization_code", "redirect_uri")},
+    ),
+    CredentialFixture(CertificateCredential, {"tenant_id": "...", "client_id": "...", "certificate_path": CERT_PATH}),
+    CredentialFixture(ClientSecretCredential, {kwarg: "..." for kwarg in ("tenant_id", "client_id", "client_secret")}),
+    CredentialFixture(DeviceCodeCredential),
+    CredentialFixture(
+        EnvironmentCredential,
+        ctor_patch=patch.dict("os.environ", {var: ".." for var in EnvironmentVariables.CLIENT_SECRET_VARS}, clear=True),
+    ),
+    CredentialFixture(InteractiveBrowserCredential),
+    CredentialFixture(UsernamePasswordCredential, {"client_id": "...", "username": "...", "password": "..."}),
+    CredentialFixture(VisualStudioCodeCredential, ctor_patch=patch(GET_USER_SETTINGS, lambda: {})),
+)
+
+all_fixtures = pytest.mark.parametrize("fixture", FIXTURES, ids=lambda fixture: fixture.cls.__name__)
+
+
+@all_fixtures
+def test_close(fixture):
+    transport = MagicMock()
+    credential = fixture(transport=transport)
+    assert not transport.__enter__.called
+    assert not transport.__exit__.called
+
+    credential.close()
+    assert not transport.__enter__.called
+    assert transport.__exit__.call_count == 1
+
+
+@all_fixtures
+def test_context_manager(fixture):
+    transport = MagicMock()
+    credential = fixture(transport=transport)
+
+    with credential:
+        assert transport.__enter__.call_count == 1
+        assert not transport.__exit__.called
+
+    assert transport.__enter__.call_count == 1
+    assert transport.__exit__.call_count == 1
+
+
+@all_fixtures
+def test_exit_args(fixture):
+    transport = MagicMock()
+    credential = fixture(transport=transport)
+    expected_args = ("type", "value", "traceback")
+    credential.__exit__(*expected_args)
+    transport.__exit__.assert_called_once_with(*expected_args)
+
+
+@pytest.mark.parametrize(
+    "cls", (AzureCliCredential, AzureApplicationCredential, AzurePowerShellCredential, EnvironmentCredential, SharedTokenCacheCredential)
+)
+def test_no_op(cls):
+    """Credentials that don't allow custom transports, or require initialization or optional config, should have no-op methods"""
+    credential = cls()
+    with credential:
+        pass
+    credential.close()

sdk/identity/azure-identity/tests/test_default.py

@@ -24,9 +24,32 @@
 from test_shared_cache_credential import build_aad_response, get_account_event, populated_cache
 
 try:
-    from unittest.mock import Mock, patch
+    from unittest.mock import MagicMock, Mock, patch
 except ImportError:  # python < 3.3
-    from mock import Mock, patch  # type: ignore
+    from mock import MagicMock, Mock, patch  # type: ignore
+
+
+def test_close():
+    transport = MagicMock()
+    credential = DefaultAzureCredential(transport=transport)
+    assert not transport.__enter__.called
+    assert not transport.__exit__.called
+
+    credential.close()
+    assert not transport.__enter__.called
+    assert transport.__exit__.called  # call count depends on the chain's composition
+
+
+def test_context_manager():
+    transport = MagicMock()
+    credential = DefaultAzureCredential(transport=transport)
+
+    with credential:
+        assert transport.__enter__.called  # call count depends on the chain's composition
+        assert not transport.__exit__.called
+
+    assert transport.__enter__.called
+    assert transport.__exit__.called
 
 
 def test_iterates_only_once():

sdk/identity/azure-identity/tests/test_environment_credential.py

@@ -9,7 +9,7 @@
 from azure.identity._constants import EnvironmentVariables
 import pytest
 
-from helpers import mock, mock_response, Request, validating_transport
+from helpers import mock
 
 
 ALL_VARIABLES = {
@@ -20,6 +20,15 @@
 }
 
 
+def test_close_incomplete_configuration():
+    EnvironmentCredential().close()
+
+
+def test_context_manager_incomplete_configuration():
+    with EnvironmentCredential():
+        pass
+
+
 def test_incomplete_configuration():
     """get_token should raise CredentialUnavailableError for incomplete configuration."""
 

sdk/identity/azure-identity/tests/test_environment_credential_async.py

@@ -12,10 +12,46 @@
 import pytest
 
 from helpers import mock_response, Request
-from helpers_async import async_validating_transport
+from helpers_async import async_validating_transport, AsyncMockTransport
 from test_environment_credential import ALL_VARIABLES
 
 
+@pytest.mark.asyncio
+async def test_close():
+    transport = AsyncMockTransport()
+    with mock.patch.dict("os.environ", {var: "..." for var in EnvironmentVariables.CLIENT_SECRET_VARS}, clear=True):
+        credential = EnvironmentCredential(transport=transport)
+    assert transport.__aexit__.call_count == 0
+
+    await credential.close()
+    assert transport.__aexit__.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_context_manager():
+    transport = AsyncMockTransport()
+    with mock.patch.dict("os.environ", {var: "..." for var in EnvironmentVariables.CLIENT_SECRET_VARS}, clear=True):
+        credential = EnvironmentCredential(transport=transport)
+
+    async with credential:
+        assert transport.__aenter__.call_count == 1
+        assert transport.__aexit__.call_count == 0
+
+    assert transport.__aenter__.call_count == 1
+    assert transport.__aexit__.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_close_incomplete_configuration():
+    await EnvironmentCredential().close()
+
+
+@pytest.mark.asyncio
+async def test_context_manager_incomplete_configuration():
+    async with EnvironmentCredential():
+        pass
+
+
 @pytest.mark.asyncio
 async def test_incomplete_configuration():
     """get_token should raise CredentialUnavailableError for incomplete configuration."""

sdk/identity/azure-identity/tests/test_managed_identity.py

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+from azure.identity._credentials import managed_identity
 import os
 import time
 
@@ -23,6 +24,51 @@
 from helpers import build_aad_response, validating_transport, mock_response, Request
 
 MANAGED_IDENTITY_ENVIRON = "azure.identity._credentials.managed_identity.os.environ"
+ALL_ENVIRONMENTS = (
+    {EnvironmentVariables.MSI_ENDPOINT: "...", EnvironmentVariables.MSI_SECRET: "..."},  # App Service
+    {EnvironmentVariables.MSI_ENDPOINT: "..."},  # Cloud Shell
+    {  # Service Fabric
+        EnvironmentVariables.IDENTITY_ENDPOINT: "...",
+        EnvironmentVariables.IDENTITY_HEADER: "...",
+        EnvironmentVariables.IDENTITY_SERVER_THUMBPRINT: "...",
+    },
+    {EnvironmentVariables.IDENTITY_ENDPOINT: "...", EnvironmentVariables.IMDS_ENDPOINT: "..."},  # Arc
+    {},  # IMDS
+)
+
+
+@pytest.mark.parametrize("environ", ALL_ENVIRONMENTS)
+def test_close(environ):
+    transport = mock.MagicMock()
+    with mock.patch.dict("os.environ", environ, clear=True):
+        credential = ManagedIdentityCredential(transport=transport)
+    assert transport.__exit__.call_count == 0
+
+    credential.close()
+    assert transport.__exit__.call_count == 1
+
+
+@pytest.mark.parametrize("environ", ALL_ENVIRONMENTS)
+def test_context_manager(environ):
+    transport = mock.MagicMock()
+    with mock.patch.dict("os.environ", environ, clear=True):
+        credential = ManagedIdentityCredential(transport=transport)
+
+    with credential:
+        assert transport.__enter__.call_count == 1
+        assert transport.__exit__.call_count == 0
+
+    assert transport.__enter__.call_count == 1
+    assert transport.__exit__.call_count == 1
+
+
+def test_close_incomplete_configuration():
+    ManagedIdentityCredential().close()
+
+
+def test_context_manager_incomplete_configuration():
+    with ManagedIdentityCredential():
+        pass
 
 
 def test_cloud_shell():
@@ -491,7 +537,9 @@ def send(request, **_):
 
     # Cloud Shell
     with mock.patch.dict(
-        MANAGED_IDENTITY_ENVIRON, {EnvironmentVariables.MSI_ENDPOINT: "https://localhost"}, clear=True,
+        MANAGED_IDENTITY_ENVIRON,
+        {EnvironmentVariables.MSI_ENDPOINT: "https://localhost"},
+        clear=True,
     ):
         credential = ManagedIdentityCredential(client_id=None, transport=mock.Mock(send=send))
         token = credential.get_token(scope)

sdk/identity/azure-identity/tests/test_managed_identity_async.py

@@ -19,19 +19,9 @@
 
 from helpers import build_aad_response, mock_response, Request
 from helpers_async import async_validating_transport, AsyncMockTransport, get_completed_future
+from test_managed_identity import ALL_ENVIRONMENTS
 
 MANAGED_IDENTITY_ENVIRON = "azure.identity.aio._credentials.managed_identity.os.environ"
-ALL_ENVIRONMENTS = (
-    {EnvironmentVariables.MSI_ENDPOINT: "...", EnvironmentVariables.MSI_SECRET: "..."},  # App Service
-    {EnvironmentVariables.MSI_ENDPOINT: "..."},  # Cloud Shell
-    {  # Service Fabric
-        EnvironmentVariables.IDENTITY_ENDPOINT: "...",
-        EnvironmentVariables.IDENTITY_HEADER: "...",
-        EnvironmentVariables.IDENTITY_SERVER_THUMBPRINT: "...",
-    },
-    {EnvironmentVariables.IDENTITY_ENDPOINT: "...", EnvironmentVariables.IMDS_ENDPOINT: "..."},  # Arc
-    {},  # IMDS
-)
 
 
 @pytest.mark.asyncio
@@ -61,6 +51,17 @@ async def test_context_manager(environ):
     assert transport.__aexit__.call_count == 1
 
 
+@pytest.mark.asyncio
+async def test_close_incomplete_configuration():
+    await ManagedIdentityCredential().close()
+
+
+@pytest.mark.asyncio
+async def test_context_manager_incomplete_configuration():
+    async with ManagedIdentityCredential():
+        pass
+
+
 @pytest.mark.asyncio
 async def test_cloud_shell():
     """Cloud Shell environment: only MSI_ENDPOINT set"""
@@ -453,7 +454,9 @@ async def send(request, **_):
     assert token.token == expected_access_token
 
     with mock.patch.dict(
-        MANAGED_IDENTITY_ENVIRON, {EnvironmentVariables.MSI_ENDPOINT: "https://localhost"}, clear=True,
+        MANAGED_IDENTITY_ENVIRON,
+        {EnvironmentVariables.MSI_ENDPOINT: "https://localhost"},
+        clear=True,
     ):
         credential = ManagedIdentityCredential(client_id=None, transport=mock.Mock(send=send))
         token = await credential.get_token(scope)