tests

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

@@ -153,8 +153,14 @@ def mock_response(status_code=200, headers=None, json_payload=None):
 
 
 def get_discovery_response(endpoint="https://a/b"):
+    """Get a mock response containing the values MSAL requires from tenant and instance discovery.
+
+    The response is incomplete and its values aren't necessarily valid, particularly for instance discovery, but it's
+    sufficient. MSAL will send token requests to "{endpoint}/oauth2/v2.0/token_endpoint" after receiving a tenant
+    discovery response created by this method.
+    """
     aad_metadata_endpoint_names = ("authorization_endpoint", "token_endpoint", "tenant_discovery_endpoint")
-    payload = {name: endpoint for name in aad_metadata_endpoint_names}
+    payload = {name: endpoint + "/oauth2/v2.0/" + name for name in aad_metadata_endpoint_names}
     payload["metadata"] = ""
     return mock_response(json_payload=payload)
 

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

@@ -56,7 +56,11 @@ def test_exceptions_do_not_expose_secrets():
 
     fns = [
         functools.partial(client.obtain_token_by_authorization_code, "code", "uri", "scope"),
-        functools.partial(client.obtain_token_by_refresh_token, "refresh token", ("scope"),),
+        functools.partial(
+            client.obtain_token_by_refresh_token,
+            "refresh token",
+            ("scope"),
+        ),
     ]
 
     def assert_secrets_not_exposed():
@@ -233,3 +237,76 @@ def test_retries_token_requests():
         client.obtain_token_by_refresh_token("", "")
     assert transport.send.call_count > 1
     transport.send.reset_mock()
+
+
+def test_shared_cache():
+    """The client should return only tokens associated with its own client_id"""
+
+    client_id_a = "client-id-a"
+    client_id_b = "client-id-b"
+    scope = "scope"
+    expected_token = "***"
+    tenant_id = "tenant"
+    authority = "https://localhost/" + tenant_id
+
+    cache = TokenCache()
+    cache.add(
+        {
+            "response": build_aad_response(access_token=expected_token),
+            "client_id": client_id_a,
+            "scope": [scope],
+            "token_endpoint": "/".join((authority, tenant_id, "oauth2/v2.0/token")),
+        }
+    )
+
+    common_args = dict(authority=authority, cache=cache, tenant_id=tenant_id)
+    client_a = AadClient(client_id=client_id_a, **common_args)
+    client_b = AadClient(client_id=client_id_b, **common_args)
+
+    # A has a cached token
+    token = client_a.get_cached_access_token([scope])
+    assert token.token == expected_token
+
+    # which B shouldn't return
+    assert client_b.get_cached_access_token([scope]) is None
+
+
+def test_multitenant_cache():
+    client_id = "client-id"
+    scope = "scope"
+    expected_token = "***"
+    tenant_a = "tenant-a"
+    tenant_b = "tenant-b"
+    tenant_c = "tenant-c"
+    authority = "https://localhost/" + tenant_a
+
+    cache = TokenCache()
+    cache.add(
+        {
+            "response": build_aad_response(access_token=expected_token),
+            "client_id": client_id,
+            "scope": [scope],
+            "token_endpoint": "/".join((authority, tenant_a, "oauth2/v2.0/token")),
+        }
+    )
+
+    common_args = dict(authority=authority, cache=cache, client_id=client_id)
+    client_a = AadClient(tenant_id=tenant_a, **common_args)
+    client_b = AadClient(tenant_id=tenant_b, **common_args)
+
+    # A has a cached token
+    token = client_a.get_cached_access_token([scope])
+    assert token.token == expected_token
+
+    # which B shouldn't return
+    assert client_b.get_cached_access_token([scope]) is None
+
+    # but C allows multitenant auth and should therefore return the token from tenant_a when appropriate
+    client_c = AadClient(tenant_id=tenant_c, allow_multitenant_authentication=True, **common_args)
+    assert client_c.get_cached_access_token([scope]) is None
+    token = client_c.get_cached_access_token([scope], tenant_id=tenant_a)
+    assert token.token == expected_token
+    with patch.dict(
+        "os.environ", {EnvironmentVariables.AZURE_IDENTITY_ENABLE_LEGACY_TENANT_SELECTION: "true"}, clear=True
+    ):
+        assert client_c.get_cached_access_token([scope], tenant_id=tenant_a) is None

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

@@ -5,13 +5,11 @@
 import functools
 from unittest.mock import Mock, patch
 from urllib.parse import urlparse
-import time
 
 from azure.core.exceptions import ClientAuthenticationError, ServiceRequestError
-from azure.identity._constants import EnvironmentVariables, DEFAULT_REFRESH_OFFSET, DEFAULT_TOKEN_REFRESH_RETRY_DELAY
+from azure.identity._constants import EnvironmentVariables
 from azure.identity._internal import AadClientCertificate
 from azure.identity.aio._internal.aad_client import AadClient
-from azure.core.credentials import AccessToken
 from msal import TokenCache
 import pytest
 
@@ -241,3 +239,76 @@ async def test_retries_token_requests():
         await client.obtain_token_by_refresh_token("", "")
     assert transport.send.call_count > 1
     transport.send.reset_mock()
+
+
+async def test_shared_cache():
+    """The client should return only tokens associated with its own client_id"""
+
+    client_id_a = "client-id-a"
+    client_id_b = "client-id-b"
+    scope = "scope"
+    expected_token = "***"
+    tenant_id = "tenant"
+    authority = "https://localhost/" + tenant_id
+
+    cache = TokenCache()
+    cache.add(
+        {
+            "response": build_aad_response(access_token=expected_token),
+            "client_id": client_id_a,
+            "scope": [scope],
+            "token_endpoint": "/".join((authority, tenant_id, "oauth2/v2.0/token")),
+        }
+    )
+
+    common_args = dict(authority=authority, cache=cache, tenant_id=tenant_id)
+    client_a = AadClient(client_id=client_id_a, **common_args)
+    client_b = AadClient(client_id=client_id_b, **common_args)
+
+    # A has a cached token
+    token = client_a.get_cached_access_token([scope])
+    assert token.token == expected_token
+
+    # which B shouldn't return
+    assert client_b.get_cached_access_token([scope]) is None
+
+
+async def test_multitenant_cache():
+    client_id = "client-id"
+    scope = "scope"
+    expected_token = "***"
+    tenant_a = "tenant-a"
+    tenant_b = "tenant-b"
+    tenant_c = "tenant-c"
+    authority = "https://localhost/" + tenant_a
+
+    cache = TokenCache()
+    cache.add(
+        {
+            "response": build_aad_response(access_token=expected_token),
+            "client_id": client_id,
+            "scope": [scope],
+            "token_endpoint": "/".join((authority, tenant_a, "oauth2/v2.0/token")),
+        }
+    )
+
+    common_args = dict(authority=authority, cache=cache, client_id=client_id)
+    client_a = AadClient(tenant_id=tenant_a, **common_args)
+    client_b = AadClient(tenant_id=tenant_b, **common_args)
+
+    # A has a cached token
+    token = client_a.get_cached_access_token([scope])
+    assert token.token == expected_token
+
+    # which B shouldn't return
+    assert client_b.get_cached_access_token([scope]) is None
+
+    # but C allows multitenant auth and should therefore return the token from tenant_a when appropriate
+    client_c = AadClient(tenant_id=tenant_c, allow_multitenant_authentication=True, **common_args)
+    assert client_c.get_cached_access_token([scope]) is None
+    token = client_c.get_cached_access_token([scope], tenant_id=tenant_a)
+    assert token.token == expected_token
+    with patch.dict(
+        "os.environ", {EnvironmentVariables.AZURE_IDENTITY_ENABLE_LEGACY_TENANT_SELECTION: "true"}, clear=True
+    ):
+        assert client_c.get_cached_access_token([scope], tenant_id=tenant_a) is None

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

@@ -2,19 +2,21 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-from azure.core.credentials import AccessToken
+from azure.core.exceptions import ClientAuthenticationError
 from azure.core.pipeline.policies import SansIOHTTPPolicy
 from azure.identity import AuthorizationCodeCredential
+from azure.identity._constants import EnvironmentVariables
 from azure.identity._internal.user_agent import USER_AGENT
 import msal
 import pytest
+from six.moves.urllib_parse import urlparse
 
 from helpers import build_aad_response, mock_response, Request, validating_transport
 
 try:
-    from unittest.mock import Mock
+    from unittest.mock import Mock, patch
 except ImportError:  # python < 3.3
-    from mock import Mock  # type: ignore
+    from mock import Mock, patch  # type: ignore
 
 
 def test_no_scopes():
@@ -114,3 +116,75 @@ def test_auth_code_credential():
     token = credential.get_token(expected_scope)
     assert token.token == expected_access_token
     assert transport.send.call_count == 2
+
+
+def test_allow_multitenant_authentication():
+    """When allow_multitenant_authentication is True, the credential should respect get_token(tenant_id=...)"""
+
+    first_tenant = "first-tenant"
+    first_token = "***"
+    second_tenant = "second-tenant"
+    second_token = first_token * 2
+
+    def send(request, **_):
+        parsed = urlparse(request.url)
+        tenant = parsed.path.split("/")[1]
+        assert tenant in (first_tenant, second_tenant), 'unexpected tenant "{}"'.format(tenant)
+        token = first_token if tenant == first_tenant else second_token
+        return mock_response(json_payload=build_aad_response(access_token=token, refresh_token="**"))
+
+    credential = AuthorizationCodeCredential(
+        first_tenant,
+        "client-id",
+        "authcode",
+        "https://localhost",
+        allow_multitenant_authentication=True,
+        transport=Mock(send=send),
+    )
+    token = credential.get_token("scope")
+    assert token.token == first_token
+
+    token = credential.get_token("scope", tenant_id=first_tenant)
+    assert token.token == first_token
+
+    token = credential.get_token("scope", tenant_id=second_tenant)
+    assert token.token == second_token
+
+    # should still default to the first tenant
+    token = credential.get_token("scope")
+    assert token.token == first_token
+
+
+def test_multitenant_authentication_not_allowed():
+    """get_token(tenant_id=...) should raise when allow_multitenant_authentication is False (the default)"""
+
+    expected_tenant = "expected-tenant"
+    expected_token = "***"
+
+    def send(request, **_):
+        parsed = urlparse(request.url)
+        tenant = parsed.path.split("/")[1]
+        token = expected_token if tenant == expected_tenant else expected_token * 2
+        return mock_response(json_payload=build_aad_response(access_token=token, refresh_token="**"))
+
+    credential = AuthorizationCodeCredential(
+        expected_tenant, "client-id", "authcode", "https://localhost", transport=Mock(send=send)
+    )
+
+    token = credential.get_token("scope")
+    assert token.token == expected_token
+
+    # explicitly specifying the configured tenant is okay
+    token = credential.get_token("scope", tenant_id=expected_tenant)
+    assert token.token == expected_token
+
+    # but any other tenant should get an error
+    with pytest.raises(ClientAuthenticationError, match="allow_multitenant_authentication"):
+        credential.get_token("scope", tenant_id="un" + expected_tenant)
+
+    # ...unless the compat switch is enabled
+    with patch.dict(
+        "os.environ", {EnvironmentVariables.AZURE_IDENTITY_ENABLE_LEGACY_TENANT_SELECTION: "true"}, clear=True
+    ):
+        token = credential.get_token("scope", tenant_id="un" + expected_tenant)
+    assert token.token == expected_token, "credential should ignore tenant_id kwarg when the compat switch is enabled"

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

@@ -2,9 +2,12 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
+from urllib.parse import urlparse
 
+from azure.core.exceptions import ClientAuthenticationError
 from azure.core.pipeline.policies import SansIOHTTPPolicy
+from azure.identity._constants import EnvironmentVariables
 from azure.identity._internal.user_agent import USER_AGENT
 from azure.identity.aio import AuthorizationCodeCredential
 import msal
@@ -137,3 +140,75 @@ async def test_auth_code_credential():
     token = await credential.get_token(expected_scope)
     assert token.token == expected_access_token
     assert transport.send.call_count == 2
+
+
+async def test_allow_multitenant_authentication():
+    """When allow_multitenant_authentication is True, the credential should respect get_token(tenant_id=...)"""
+
+    first_tenant = "first-tenant"
+    first_token = "***"
+    second_tenant = "second-tenant"
+    second_token = first_token * 2
+
+    async def send(request, **_):
+        parsed = urlparse(request.url)
+        tenant = parsed.path.split("/")[1]
+        assert tenant in (first_tenant, second_tenant), 'unexpected tenant "{}"'.format(tenant)
+        token = first_token if tenant == first_tenant else second_token
+        return mock_response(json_payload=build_aad_response(access_token=token, refresh_token="**"))
+
+    credential = AuthorizationCodeCredential(
+        first_tenant,
+        "client-id",
+        "authcode",
+        "https://localhost",
+        allow_multitenant_authentication=True,
+        transport=Mock(send=send),
+    )
+    token = await credential.get_token("scope")
+    assert token.token == first_token
+
+    token = await credential.get_token("scope", tenant_id=first_tenant)
+    assert token.token == first_token
+
+    token = await credential.get_token("scope", tenant_id=second_tenant)
+    assert token.token == second_token
+
+    # should still default to the first tenant
+    token = await credential.get_token("scope")
+    assert token.token == first_token
+
+
+async def test_multitenant_authentication_not_allowed():
+    """get_token(tenant_id=...) should raise when allow_multitenant_authentication is False (the default)"""
+
+    expected_tenant = "expected-tenant"
+    expected_token = "***"
+
+    async def send(request, **_):
+        parsed = urlparse(request.url)
+        tenant = parsed.path.split("/")[1]
+        token = expected_token if tenant == expected_tenant else expected_token * 2
+        return mock_response(json_payload=build_aad_response(access_token=token, refresh_token="**"))
+
+    credential = AuthorizationCodeCredential(
+        expected_tenant, "client-id", "authcode", "https://localhost", transport=Mock(send=send)
+    )
+
+    token = await credential.get_token("scope")
+    assert token.token == expected_token
+
+    # explicitly specifying the configured tenant is okay
+    token = await credential.get_token("scope", tenant_id=expected_tenant)
+    assert token.token == expected_token
+
+    # but any other tenant should get an error
+    with pytest.raises(ClientAuthenticationError, match="allow_multitenant_authentication"):
+        await credential.get_token("scope", tenant_id="un" + expected_tenant)
+
+    # ...unless the compat switch is enabled
+    with patch.dict(
+        "os.environ", {EnvironmentVariables.AZURE_IDENTITY_ENABLE_LEGACY_TENANT_SELECTION: "true"}, clear=True
+    ):
+        token = await credential.get_token("scope", tenant_id="un" + expected_tenant)
+    assert token.token == expected_token, "credential should ignore tenant_id kwarg when the compat switch is enabled"