[Core] Add refresh_on property to AccessToken (#36183)

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

Author: pvaneck

sdk/core/azure-core/CHANGELOG.md

@@ -1,9 +1,12 @@
 # Release History
 
-## 1.30.3 (Unreleased)
+## 1.31.0 (Unreleased)
 
 ### Features Added
 
+- `AccessToken` now has an optional `refresh_on` attribute that can be used to specify when the token should be refreshed.  #36183
+  - `BearerTokenCredentialPolicy` and `AsyncBearerTokenCredentialPolicy` now check the `refresh_on` attribute when determining if a token request should be made.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/core/azure-core/azure/core/_version.py

@@ -9,4 +9,4 @@
 # regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "1.30.3"
+VERSION = "1.31.0"

sdk/core/azure-core/azure/core/credentials.py

@@ -12,10 +12,12 @@ class AccessToken(NamedTuple):
 
     token: str
     expires_on: int
+    refresh_on: Optional[int] = None
 
 
 AccessToken.token.__doc__ = """The token string."""
 AccessToken.expires_on.__doc__ = """The token's expiration time in Unix time."""
+AccessToken.refresh_on.__doc__ = """When the token should be refreshed in Unix time."""
 
 
 @runtime_checkable

sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py

@@ -69,7 +69,12 @@ def _update_headers(headers: MutableMapping[str, str], token: str) -> None:
 
     @property
     def _need_new_token(self) -> bool:
-        return not self._token or self._token.expires_on - time.time() < 300
+        now = time.time()
+        return (
+            not self._token
+            or (self._token.refresh_on is not None and self._token.refresh_on <= now)
+            or self._token.expires_on - now < 300
+        )
 
 
 class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, HTTPPolicy[HTTPRequestType, HTTPResponseType]):

sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py

@@ -164,4 +164,9 @@ def on_exception(self, request: PipelineRequest[HTTPRequestType]) -> None:
         return
 
     def _need_new_token(self) -> bool:
-        return not self._token or self._token.expires_on - time.time() < 300
+        now = time.time()
+        return (
+            not self._token
+            or (self._token.refresh_on is not None and self._token.refresh_on <= now)
+            or self._token.expires_on - now < 300
+        )

sdk/core/azure-core/tests/test_authentication.py

@@ -305,6 +305,33 @@ def test_key_vault_regression(http_request):
     assert policy._token.token == token
 
 
+def test_need_new_token():
+    expected_scope = "scope"
+    now = int(time.time())
+
+    policy = BearerTokenCredentialPolicy(Mock(), expected_scope)
+
+    # Token is expired.
+    policy._token = AccessToken("", now - 1200)
+    assert policy._need_new_token
+
+    # Token is about to expire within 300 seconds.
+    policy._token = AccessToken("", now + 299)
+    assert policy._need_new_token
+
+    # Token still has more than 300 seconds to live.
+    policy._token = AccessToken("", now + 305)
+    assert not policy._need_new_token
+
+    # Token has both expires_on and refresh_on set well into the future.
+    policy._token = AccessToken("", now + 1200, now + 1200)
+    assert not policy._need_new_token
+
+    # Token is not close to expiring, but refresh_on is in the past.
+    policy._token = AccessToken("", now + 1200, now - 1)
+    assert policy._need_new_token
+
+
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 def test_azure_key_credential_policy(http_request):
     """Tests to see if we can create an AzureKeyCredentialPolicy"""