Fix ManagedIdentityCredential token caching (Azure#17323)

Author: chlowell

sdk/identity/azure-identity/CHANGELOG.md

@@ -1,7 +1,8 @@
 # Release History
 
 ## 1.6.0b3 (Unreleased)
-
+### Fixed
+- ManagedIdentityCredential caches tokens correctly
 
 ## 1.6.0b2 (2021-03-09)
 ### Breaking Changes

sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py

@@ -79,7 +79,7 @@ def _process_response(self, response, request_time):
 
         # caching is the final step because TokenCache.add mutates its "event"
         self._cache.add(
-            event={"response": content, "scope": content["resource"]}, now=request_time,
+            event={"response": content, "scope": [content["resource"]]}, now=request_time,
         )
 
         return token
@@ -89,8 +89,9 @@ def get_cached_token(self, *scopes):
         resource = _scopes_to_resource(*scopes)
         tokens = self._cache.find(TokenCache.CredentialType.ACCESS_TOKEN, target=[resource])
         for token in tokens:
-            if token["expires_on"] > time.time():
-                return AccessToken(token["secret"], token["expires_on"])
+            expires_on = int(token["expires_on"])
+            if expires_on > time.time():
+                return AccessToken(token["secret"], expires_on)
         return None
 
     @abc.abstractmethod

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

@@ -0,0 +1,44 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import time
+
+from azure.core.pipeline.transport import HttpRequest
+from azure.identity._internal.managed_identity_client import ManagedIdentityClient
+
+from helpers import mock_response, Request, validating_transport
+
+
+def test_caching():
+    scope = "scope"
+    expected_expires_on = int(time.time() + 3600)
+    expected_token = "*"
+    transport = validating_transport(
+        requests=[Request(url="http://localhost")],
+        responses=[
+            mock_response(
+                json_payload={
+                    "access_token": expected_token,
+                    "expires_in": 3600,
+                    "expires_on": expected_expires_on,
+                    "resource": scope,
+                    "token_type": "Bearer",
+                }
+            )
+        ],
+    )
+    client = ManagedIdentityClient(
+        request_factory=lambda _, __: HttpRequest("GET", "http://localhost"), transport=transport
+    )
+
+    token = client.get_cached_token(scope)
+    assert not token
+
+    token = client.request_token(scope)
+    assert token.expires_on == expected_expires_on
+    assert token.token == expected_token
+
+    token = client.get_cached_token(scope)
+    assert token.expires_on == expected_expires_on
+    assert token.token == expected_token

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

@@ -0,0 +1,47 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import time
+
+from azure.core.pipeline.transport import HttpRequest
+from azure.identity.aio._internal.managed_identity_client import AsyncManagedIdentityClient
+import pytest
+
+from helpers import mock_response, Request
+from helpers_async import async_validating_transport
+
+
+@pytest.mark.asyncio
+async def test_caching():
+    scope = "scope"
+    expected_expires_on = int(time.time() + 3600)
+    expected_token = "*"
+    transport = async_validating_transport(
+        requests=[Request(url="http://localhost")],
+        responses=[
+            mock_response(
+                json_payload={
+                    "access_token": expected_token,
+                    "expires_in": 3600,
+                    "expires_on": expected_expires_on,
+                    "resource": scope,
+                    "token_type": "Bearer",
+                }
+            )
+        ],
+    )
+    client = AsyncManagedIdentityClient(
+        request_factory=lambda _, __: HttpRequest("GET", "http://localhost"), transport=transport
+    )
+
+    token = client.get_cached_token(scope)
+    assert not token
+
+    token = await client.request_token(scope)
+    assert token.expires_on == expected_expires_on
+    assert token.token == expected_token
+
+    token = client.get_cached_token(scope)
+    assert token.expires_on == expected_expires_on
+    assert token.token == expected_token