[Identity] Fix issue with cae_cache not being used (#42145)

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/identity/azure-identity/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed an issue where CAE (Continuous Access Evaluation) caches were not properly used by `AuthorizationCodeCredential` and the asynchronous `OnBehalfOfCredential`. ([#42145](https://github.com/Azure/azure-sdk-for-python/pull/42145))
+
 ### Other Changes
 
 ## 1.24.0b1 (2025-07-17)

sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py

@@ -127,7 +127,7 @@ def _request_token(self, *scopes: str, **kwargs) -> AccessTokenInfo:
             return token
 
         token = None
-        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
+        for refresh_token in self._client.get_cached_refresh_tokens(scopes, **kwargs):
             if "secret" in refresh_token:
                 token = self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
                 if token:

sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py

@@ -133,7 +133,7 @@ async def _request_token(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo:
             return token
 
         token = cast(AccessTokenInfo, None)
-        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
+        for refresh_token in self._client.get_cached_refresh_tokens(scopes, **kwargs):
             if "secret" in refresh_token:
                 token = await self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
                 if token:

sdk/identity/azure-identity/azure/identity/aio/_credentials/on_behalf_of.py

@@ -119,7 +119,7 @@ async def _request_token(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo:
         # Note we assume the cache has tokens for one user only. That's okay because each instance of this class is
         # locked to a single user (assertion). This assumption will become unsafe if this class allows applications
         # to change an instance's assertion.
-        refresh_tokens = self._client.get_cached_refresh_tokens(scopes)
+        refresh_tokens = self._client.get_cached_refresh_tokens(scopes, **kwargs)
         if len(refresh_tokens) == 1:  # there should be only one
             try:
                 refresh_token = refresh_tokens[0]["secret"]

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

@@ -80,7 +80,8 @@ def test_tenant_id(get_token_method):
 
 
 @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
-def test_auth_code_credential(get_token_method):
+@pytest.mark.parametrize("enable_cae", [True, False])
+def test_auth_code_credential(get_token_method, enable_cae):
     client_id = "client id"
     secret = "fake-client-secret"
     tenant_id = "tenant"
@@ -118,6 +119,7 @@ def test_auth_code_credential(get_token_method):
         responses=[mock_response(json_payload=auth_response)] * 2,
     )
     cache = msal.TokenCache()
+    cae_cache = msal.TokenCache()
 
     credential = AuthorizationCodeCredential(
         client_id=client_id,
@@ -127,22 +129,29 @@ def test_auth_code_credential(get_token_method):
         redirect_uri=redirect_uri,
         transport=transport,
         cache=cache,
+        cae_cache=cae_cache,
     )
 
     # first call should redeem the auth code
-    token = getattr(credential, get_token_method)(expected_scope)
+    kwargs = {"enable_cae": enable_cae}
+    if get_token_method == "get_token_info":
+        kwargs = {"options": kwargs}
+    token = getattr(credential, get_token_method)(expected_scope, **kwargs)
     assert token.token == expected_access_token
     assert transport.send.call_count == 1
 
     # no auth code -> credential should return cached token
-    token = getattr(credential, get_token_method)(expected_scope)
+    token = getattr(credential, get_token_method)(expected_scope, **kwargs)
     assert token.token == expected_access_token
     assert transport.send.call_count == 1
 
     # no auth code, no cached token -> credential should redeem refresh token
-    cached_access_token = list(cache.search(cache.CredentialType.ACCESS_TOKEN))[0]
-    cache.remove_at(cached_access_token)
-    token = getattr(credential, get_token_method)(expected_scope)
+    cache_being_used = cae_cache if enable_cae else cache
+    cached_tokens = list(cache_being_used.search(cache_being_used.CredentialType.ACCESS_TOKEN))
+    assert cached_tokens
+    cached_access_token = cached_tokens[0]
+    cache_being_used.remove_at(cached_access_token)
+    token = getattr(credential, get_token_method)(expected_scope, **kwargs)
     assert token.token == expected_access_token
     assert transport.send.call_count == 2
 

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

@@ -249,7 +249,8 @@ def test_invalid_cert():
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
-async def test_refresh_token(get_token_method):
+@pytest.mark.parametrize("enable_cae", [True, False])
+async def test_refresh_token(get_token_method, enable_cae):
     first_token = "***"
     second_token = first_token * 2
     refresh_token = "refresh-token"
@@ -274,10 +275,15 @@ async def send(request, **kwargs):
     credential = OnBehalfOfCredential(
         "tenant-id", "client-id", client_secret="secret", user_assertion="assertion", transport=Mock(send=send)
     )
-    token = await getattr(credential, get_token_method)("scope")
+
+    kwargs = {"enable_cae": enable_cae}
+    if get_token_method == "get_token_info":
+        kwargs = {"options": kwargs}
+
+    token = await getattr(credential, get_token_method)("scope", **kwargs)
     assert token.token == first_token
 
-    token = await getattr(credential, get_token_method)("scope")
+    token = await getattr(credential, get_token_method)("scope", **kwargs)
     assert token.token == second_token
 
     assert requests == 2