From c40fe7a4448a296cede2d4d9d77f0ffb755c70ee Mon Sep 17 00:00:00 2001
From: Charles Lowell <chlowe@microsoft.com>
Date: Wed, 31 Mar 2021 14:49:07 -0700
Subject: [PATCH 1/2] consolidate proactive refresh logic

---
 .../azure/identity/_authn_client.py           | 12 ------
 .../_credentials/authorization_code.py        | 39 ++++++++-----------
 .../identity/_credentials/managed_identity.py | 29 ++++----------
 .../azure/identity/_credentials/vscode.py     | 32 ++++++---------
 .../identity/_internal/aad_client_base.py     | 13 -------
 .../aio/_credentials/authorization_code.py    | 38 +++++++-----------
 .../identity/aio/_credentials/certificate.py  | 36 +++++------------
 .../aio/_credentials/client_secret.py         | 39 +++++--------------
 .../aio/_credentials/managed_identity.py      | 31 +++------------
 .../azure/identity/aio/_credentials/vscode.py | 28 +++++--------
 10 files changed, 82 insertions(+), 215 deletions(-)

diff --git a/sdk/identity/azure-identity/azure/identity/_authn_client.py b/sdk/identity/azure-identity/azure/identity/_authn_client.py
index 5c2801abd5eb..23901b0cae3d 100644
--- a/sdk/identity/azure-identity/azure/identity/_authn_client.py
+++ b/sdk/identity/azure-identity/azure/identity/_authn_client.py
@@ -75,18 +75,6 @@ def __init__(self, endpoint=None, authority=None, tenant=None, **kwargs):  # pyl
     def auth_url(self):
         return self._auth_url
 
-    def should_refresh(self, token):
-        # type: (AccessToken) -> bool
-        """ check if the token needs refresh or not
-        """
-        expires_on = int(token.expires_on)
-        now = int(time.time())
-        if expires_on - now > self._token_refresh_offset:
-            return False
-        if now - self._last_refresh_time < self._token_refresh_retry_delay:
-            return False
-        return True
-
     def get_cached_token(self, scopes):
         # type: (Iterable[str]) -> Optional[AccessToken]
         tokens = self._cache.find(TokenCache.CredentialType.ACCESS_TOKEN, target=list(scopes))
diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py b/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py
index 204c895acab5..0017991e70d4 100644
--- a/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py
+++ b/sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py
@@ -6,7 +6,7 @@
 
 from azure.core.exceptions import ClientAuthenticationError
 from .._internal.aad_client import AadClient
-from .._internal.decorators import log_get_token
+from .._internal.get_token_mixin import GetTokenMixin
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
@@ -14,7 +14,7 @@
     from azure.core.credentials import AccessToken
 
 
-class AuthorizationCodeCredential(object):
+class AuthorizationCodeCredential(GetTokenMixin):
     """Authenticates by redeeming an authorization code previously obtained from Azure Active Directory.
 
     See https://docs.microsoft.com/azure/active-directory/develop/v2-oauth2-auth-code-flow for more information
@@ -38,8 +38,8 @@ def __init__(self, tenant_id, client_id, authorization_code, redirect_uri, **kwa
         self._client_secret = kwargs.pop("client_secret", None)
         self._client = kwargs.pop("client", None) or AadClient(tenant_id, client_id, **kwargs)
         self._redirect_uri = redirect_uri
+        super(AuthorizationCodeCredential, self).__init__()
 
-    @log_get_token("AuthorizationCodeCredential")
     def get_token(self, *scopes, **kwargs):
         # type: (*str, **Any) -> AccessToken
         """Request an access token for `scopes`.
@@ -56,9 +56,14 @@ def get_token(self, *scopes, **kwargs):
           attribute gives a reason. Any error response from Azure Active Directory is available as the error's
           ``response`` attribute.
         """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
+        return super(AuthorizationCodeCredential, self).get_token(*scopes)
 
+    def _acquire_token_silently(self, *scopes):
+        # type: (*str) -> Optional[AccessToken]
+        return self._client.get_cached_access_token(scopes)
+
+    def _request_token(self, *scopes, **kwargs):
+        # type: (*str, **Any) -> AccessToken
         if self._authorization_code:
             token = self._client.obtain_token_by_authorization_code(
                 scopes=scopes, code=self._authorization_code, redirect_uri=self._redirect_uri, **kwargs
@@ -66,14 +71,12 @@ def get_token(self, *scopes, **kwargs):
             self._authorization_code = None  # auth codes are single-use
             return token
 
-        token = self._client.get_cached_access_token(scopes)
-        if not token:
-            token = self._redeem_refresh_token(scopes, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                self._redeem_refresh_token(scopes, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
+        token = None
+        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
+            if "secret" in refresh_token:
+                token = self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
+                if token:
+                    break
 
         if not token:
             raise ClientAuthenticationError(
@@ -81,13 +84,3 @@ def get_token(self, *scopes, **kwargs):
             )
 
         return token
-
-    def _redeem_refresh_token(self, scopes, **kwargs):
-        # type: (Iterable[str], **Any) -> Optional[AccessToken]
-        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
-            if "secret" not in refresh_token:
-                continue
-            token = self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
-            if token:
-                return token
-        return None
diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py
index e2a231571216..d2001a1002eb 100644
--- a/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py
+++ b/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py
@@ -23,6 +23,7 @@
 from .._authn_client import AuthnClient
 from .._constants import Endpoints, EnvironmentVariables
 from .._internal.decorators import log_get_token
+from .._internal.get_token_mixin import GetTokenMixin
 from .._internal.user_agent import USER_AGENT
 
 try:
@@ -156,7 +157,7 @@ def _create_config(**kwargs):
     }
 
 
-class ImdsCredential(_ManagedIdentityBase):
+class ImdsCredential(_ManagedIdentityBase, GetTokenMixin):
     """Authenticates with a managed identity via the IMDS endpoint.
 
     :keyword str client_id: ID of a user-assigned identity. Leave unspecified to use a system-assigned identity.
@@ -167,16 +168,12 @@ def __init__(self, **kwargs):
         super(ImdsCredential, self).__init__(endpoint=Endpoints.IMDS, client_cls=AuthnClient, **kwargs)
         self._endpoint_available = None  # type: Optional[bool]
 
-    def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
-        # type: (*str, **Any) -> AccessToken
-        """Request an access token for `scopes`.
-
-        This method is called automatically by Azure SDK clients.
+    def _acquire_token_silently(self, *scopes):
+        # type: (*str) -> Optional[AccessToken]
+        return self._client.get_cached_token(scopes)
 
-        :param str scopes: desired scope for the access token. This credential allows only one scope per request.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.identity.CredentialUnavailableError: the IMDS endpoint is unreachable
-        """
+    def _request_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
+        # type: (*str, **Any) -> AccessToken
         if self._endpoint_available is None:
             # Lacking another way to determine whether the IMDS endpoint is listening,
             # we send a request it would immediately reject (missing a required header),
@@ -199,18 +196,6 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
         if len(scopes) != 1:
             raise ValueError("This credential requires exactly one scope per token request.")
 
-        token = self._client.get_cached_token(scopes)
-        if not token:
-            token = self._refresh_token(*scopes)
-        elif self._client.should_refresh(token):
-            try:
-                token = self._refresh_token(*scopes)
-            except Exception:  # pylint: disable=broad-except
-                pass
-
-        return token
-
-    def _refresh_token(self, *scopes):
         resource = scopes[0]
         if resource.endswith("/.default"):
             resource = resource[: -len("/.default")]
diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/vscode.py b/sdk/identity/azure-identity/azure/identity/_credentials/vscode.py
index 6eb38455a62a..8e07917ef5f0 100644
--- a/sdk/identity/azure-identity/azure/identity/_credentials/vscode.py
+++ b/sdk/identity/azure-identity/azure/identity/_credentials/vscode.py
@@ -9,7 +9,7 @@
 from .._constants import AZURE_VSCODE_CLIENT_ID
 from .._internal import validate_tenant_id
 from .._internal.aad_client import AadClient
-from .._internal.decorators import log_get_token
+from .._internal.get_token_mixin import GetTokenMixin
 
 if sys.platform.startswith("win"):
     from .._internal.win_vscode_adapter import get_credentials
@@ -20,11 +20,11 @@
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
-    from typing import Any, Iterable, Optional
+    from typing import Any, Optional
     from azure.core.credentials import AccessToken
 
 
-class VisualStudioCodeCredential(object):
+class VisualStudioCodeCredential(GetTokenMixin):
     """Authenticates as the Azure user signed in to Visual Studio Code.
 
     :keyword str authority: Authority of an Azure Active Directory endpoint, for example 'login.microsoftonline.com',
@@ -36,6 +36,7 @@ class VisualStudioCodeCredential(object):
 
     def __init__(self, **kwargs):
         # type: (**Any) -> None
+        super(VisualStudioCodeCredential, self).__init__()
         self._refresh_token = None
         self._client = kwargs.pop("_client", None)
         self._tenant_id = kwargs.pop("tenant_id", None) or "organizations"
@@ -43,7 +44,6 @@ def __init__(self, **kwargs):
         if not self._client:
             self._client = AadClient(self._tenant_id, AZURE_VSCODE_CLIENT_ID, **kwargs)
 
-    @log_get_token("VisualStudioCodeCredential")
     def get_token(self, *scopes, **kwargs):
         # type: (*str, **Any) -> AccessToken
         """Request an access token for `scopes` as the user currently signed in to Visual Studio Code.
@@ -55,31 +55,21 @@ def get_token(self, *scopes, **kwargs):
         :raises ~azure.identity.CredentialUnavailableError: the credential cannot retrieve user details from Visual
           Studio Code
         """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
-
         if self._tenant_id.lower() == "adfs":
             raise CredentialUnavailableError(
                 message="VisualStudioCodeCredential authentication unavailable. ADFS is not supported."
             )
+        return super(VisualStudioCodeCredential, self).get_token(*scopes, **kwargs)
 
-        token = self._client.get_cached_access_token(scopes)
-
-        if not token:
-            token = self._redeem_refresh_token(scopes, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                self._redeem_refresh_token(scopes, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
+    def _acquire_token_silently(self, *scopes):
+        # type: (*str) -> Optional[AccessToken]
+        return self._client.get_cached_access_token(scopes)
 
-    def _redeem_refresh_token(self, scopes, **kwargs):
-        # type: (Iterable[str], **Any) -> Optional[AccessToken]
+    def _request_token(self, *scopes, **kwargs):
+        # type: (*str, **Any) -> AccessToken
         if not self._refresh_token:
             self._refresh_token = get_credentials()
             if not self._refresh_token:
                 raise CredentialUnavailableError(message="Failed to get Azure user details from Visual Studio Code.")
 
-        token = self._client.obtain_token_by_refresh_token(scopes, self._refresh_token, **kwargs)
-        return token
+        return self._client.obtain_token_by_refresh_token(scopes, self._refresh_token, **kwargs)
diff --git a/sdk/identity/azure-identity/azure/identity/_internal/aad_client_base.py b/sdk/identity/azure-identity/azure/identity/_internal/aad_client_base.py
index 3e13b7739835..511b05982a72 100644
--- a/sdk/identity/azure-identity/azure/identity/_internal/aad_client_base.py
+++ b/sdk/identity/azure-identity/azure/identity/_internal/aad_client_base.py
@@ -69,19 +69,6 @@ def get_cached_refresh_tokens(self, scopes):
         """Assumes all cached refresh tokens belong to the same user"""
         return self._cache.find(TokenCache.CredentialType.REFRESH_TOKEN, target=list(scopes))
 
-    def should_refresh(self, token):
-        # type: (AccessToken) -> bool
-        """ check if the token needs refresh or not
-        """
-        expires_on = int(token.expires_on)
-        now = int(time.time())
-        if expires_on - now > self._token_refresh_offset:
-            return False
-        if now - self._last_refresh_time < self._token_refresh_retry_delay:
-            return False
-        return True
-
-
     @abc.abstractmethod
     def obtain_token_by_authorization_code(self, scopes, code, redirect_uri, client_secret=None, **kwargs):
         pass
diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py
index 9a55f19327b6..d80f59a9e973 100644
--- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py
+++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py
@@ -6,7 +6,7 @@
 
 from azure.core.exceptions import ClientAuthenticationError
 from .._internal import AadClient, AsyncContextManager
-from .._internal.decorators import log_get_token_async
+from .._internal.get_token_mixin import GetTokenMixin
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
@@ -14,7 +14,7 @@
     from azure.core.credentials import AccessToken
 
 
-class AuthorizationCodeCredential(AsyncContextManager):
+class AuthorizationCodeCredential(AsyncContextManager, GetTokenMixin):
     """Authenticates by redeeming an authorization code previously obtained from Azure Active Directory.
 
     See https://docs.microsoft.com/azure/active-directory/develop/v2-oauth2-auth-code-flow for more information
@@ -50,8 +50,8 @@ def __init__(
         self._client_secret = kwargs.pop("client_secret", None)
         self._client = kwargs.pop("client", None) or AadClient(tenant_id, client_id, **kwargs)
         self._redirect_uri = redirect_uri
+        super().__init__()
 
-    @log_get_token_async
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         """Request an access token for `scopes`.
 
@@ -67,37 +67,29 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
           attribute gives a reason. Any error response from Azure Active Directory is available as the error's
           ``response`` attribute.
         """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
+        return await super().get_token(*scopes, **kwargs)
 
+    async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]":
+        return self._client.get_cached_access_token(scopes)
+
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         if self._authorization_code:
             token = await self._client.obtain_token_by_authorization_code(
                 scopes=scopes, code=self._authorization_code, redirect_uri=self._redirect_uri, **kwargs
             )
-
             self._authorization_code = None  # auth codes are single-use
             return token
 
-        token = self._client.get_cached_access_token(scopes)
-        if not token:
-            token = await self._redeem_refresh_token(scopes, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                await self._redeem_refresh_token(scopes, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
+        token = None
+        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
+            if "secret" in refresh_token:
+                token = await self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
+                if token:
+                    break
+
         if not token:
             raise ClientAuthenticationError(
                 message="No authorization code, cached access token, or refresh token available."
             )
 
         return token
-
-    async def _redeem_refresh_token(self, scopes: "Iterable[str]", **kwargs: "Any") -> "Optional[AccessToken]":
-        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
-            if "secret" not in refresh_token:
-                continue
-            token = await self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
-            if token:
-                return token
-        return None
diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/certificate.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/certificate.py
index 97b6996ced6c..1de0bb2baaa0 100644
--- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/certificate.py
+++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/certificate.py
@@ -7,17 +7,17 @@
 import msal
 
 from .._internal import AadClient, AsyncContextManager
-from .._internal.decorators import log_get_token_async
+from .._internal.get_token_mixin import GetTokenMixin
 from ..._credentials.certificate import get_client_credential
 from ..._internal import AadClientCertificate, validate_tenant_id
-from ..._persistent_cache import _load_persistent_cache, TokenCachePersistenceOptions
+from ..._persistent_cache import _load_persistent_cache
 
 if TYPE_CHECKING:
     from typing import Any, Optional
     from azure.core.credentials import AccessToken
 
 
-class CertificateCredential(AsyncContextManager):
+class CertificateCredential(AsyncContextManager, GetTokenMixin):
     """Authenticates as a service principal using a certificate.
 
     The certificate must have an RSA private key, because this credential signs assertions using RS256.
@@ -57,6 +57,7 @@ def __init__(self, tenant_id, client_id, certificate_path=None, **kwargs):
 
         self._client = AadClient(tenant_id, client_id, cache=cache, **kwargs)
         self._client_id = client_id
+        super().__init__()
 
     async def __aenter__(self):
         await self._client.__aenter__()
@@ -67,27 +68,8 @@ async def close(self):
 
         await self._client.__aexit__()
 
-    @log_get_token_async
-    async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
-        """Asynchronously request an access token for `scopes`.
-
-        This method is called automatically by Azure SDK clients.
-
-        :param str scopes: desired scopes for the access token. This method requires at least one scope.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message``
-          attribute gives a reason. Any error response from Azure Active Directory is available as the error's
-          ``response`` attribute.
-        """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
-
-        token = self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
-        if not token:
-            token = await self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                await self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
+    async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]":
+        return self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
+
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        return await self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/client_secret.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/client_secret.py
index 335e61c989cb..7bceaae674db 100644
--- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/client_secret.py
+++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/client_secret.py
@@ -7,16 +7,16 @@
 import msal
 
 from .._internal import AadClient, AsyncContextManager
-from .._internal.decorators import log_get_token_async
+from .._internal.get_token_mixin import GetTokenMixin
 from ..._internal import validate_tenant_id
-from ..._persistent_cache import _load_persistent_cache, TokenCachePersistenceOptions
+from ..._persistent_cache import _load_persistent_cache
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Optional
     from azure.core.credentials import AccessToken
 
 
-class ClientSecretCredential(AsyncContextManager):
+class ClientSecretCredential(AsyncContextManager, GetTokenMixin):
     """Authenticates as a service principal using a client ID and client secret.
 
     :param str tenant_id: ID of the service principal's tenant. Also called its 'directory' ID.
@@ -31,8 +31,7 @@ class ClientSecretCredential(AsyncContextManager):
     :paramtype cache_persistence_options: ~azure.identity.TokenCachePersistenceOptions
     """
 
-    def __init__(self, tenant_id, client_id, client_secret, **kwargs):
-        # type: (str, str, str, **Any) -> None
+    def __init__(self, tenant_id: str, client_id: str, client_secret: str, **kwargs: "Any") -> None:
         if not client_id:
             raise ValueError("client_id should be the id of an Azure Active Directory application")
         if not client_secret:
@@ -52,6 +51,7 @@ def __init__(self, tenant_id, client_id, client_secret, **kwargs):
         self._client = AadClient(tenant_id, client_id, cache=cache, **kwargs)
         self._client_id = client_id
         self._secret = client_secret
+        super().__init__()
 
     async def __aenter__(self):
         await self._client.__aenter__()
@@ -62,27 +62,8 @@ async def close(self):
 
         await self._client.__aexit__()
 
-    @log_get_token_async
-    async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
-        """Asynchronously request an access token for `scopes`.
+    async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]":
+        return self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
 
-        This method is called automatically by Azure SDK clients.
-
-        :param str scopes: desired scopes for the access token. This method requires at least one scope.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message``
-          attribute gives a reason. Any error response from Azure Active Directory is available as the error's
-          ``response`` attribute.
-        """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
-
-        token = self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
-        if not token:
-            token = await self._client.obtain_token_by_client_secret(scopes, self._secret, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                await self._client.obtain_token_by_client_secret(scopes, self._secret, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        return await self._client.obtain_token_by_client_secret(scopes, self._secret, **kwargs)
diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py
index 30444eaa41b4..d2d64f80ee6e 100644
--- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py
+++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py
@@ -2,7 +2,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import abc
 import logging
 import os
 from typing import TYPE_CHECKING
@@ -14,6 +13,7 @@
 from .._authn_client import AsyncAuthnClient
 from .._internal import AsyncContextManager
 from .._internal.decorators import log_get_token_async
+from .._internal.get_token_mixin import GetTokenMixin
 from ... import CredentialUnavailableError
 from ..._constants import Endpoints, EnvironmentVariables
 from ..._credentials.managed_identity import _ManagedIdentityBase
@@ -109,17 +109,13 @@ async def close(self):
 
         await self._client.__aexit__()
 
-    @abc.abstractmethod
-    async def get_token(self, *scopes, **kwargs):
-        pass
-
     @staticmethod
     def _create_config(**kwargs: "Any") -> "Configuration":
         """Build a default configuration for the credential's HTTP pipeline."""
         return _ManagedIdentityBase._create_config(retry_policy=AsyncRetryPolicy, **kwargs)
 
 
-class ImdsCredential(_AsyncManagedIdentityBase):
+class ImdsCredential(_AsyncManagedIdentityBase, GetTokenMixin):
     """Asynchronously authenticates with a managed identity via the IMDS endpoint.
 
     :keyword str client_id: ID of a user-assigned identity. Leave unspecified to use a system-assigned identity.
@@ -129,15 +125,10 @@ def __init__(self, **kwargs: "Any") -> None:
         super().__init__(endpoint=Endpoints.IMDS, **kwargs)
         self._endpoint_available = None  # type: Optional[bool]
 
-    async def get_token(self, *scopes: str, **kwargs: "Any") -> AccessToken:  # pylint:disable=unused-argument
-        """Asynchronously request an access token for `scopes`.
-
-        This method is called automatically by Azure SDK clients.
+    async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]":
+        return self._client.get_cached_token(scopes)
 
-        :param str scopes: desired scope for the access token. This credential allows only one scope per request.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.identity.CredentialUnavailableError: the IMDS endpoint is unreachable
-        """
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
         if self._endpoint_available is None:
             # Lacking another way to determine whether the IMDS endpoint is listening,
             # we send a request it would immediately reject (missing a required header),
@@ -160,18 +151,6 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> AccessToken:  # pyli
         if len(scopes) != 1:
             raise ValueError("This credential requires exactly one scope per token request.")
 
-        token = self._client.get_cached_token(scopes)
-        if not token:
-            token = await self._refresh_token(*scopes)
-        elif self._client.should_refresh(token):
-            try:
-                token = await self._refresh_token(*scopes)
-            except Exception:  # pylint: disable=broad-except
-                pass
-
-        return token
-
-    async def _refresh_token(self, *scopes):
         resource = scopes[0]
         if resource.endswith("/.default"):
             resource = resource[: -len("/.default")]
diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/vscode.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/vscode.py
index dab20b06bb65..c13e239e385e 100644
--- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/vscode.py
+++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/vscode.py
@@ -8,17 +8,17 @@
 from ..._constants import AZURE_VSCODE_CLIENT_ID
 from .._internal import AsyncContextManager
 from .._internal.aad_client import AadClient
-from .._internal.decorators import log_get_token_async
+from .._internal.get_token_mixin import GetTokenMixin
 from ..._credentials.vscode import get_credentials
 from ..._internal import validate_tenant_id
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
-    from typing import Any, Iterable, Optional
+    from typing import Any, Optional
     from azure.core.credentials import AccessToken
 
 
-class VisualStudioCodeCredential(AsyncContextManager):
+class VisualStudioCodeCredential(AsyncContextManager, GetTokenMixin):
     """Authenticates as the Azure user signed in to Visual Studio Code.
 
     :keyword str authority: Authority of an Azure Active Directory endpoint, for example 'login.microsoftonline.com',
@@ -29,6 +29,7 @@ class VisualStudioCodeCredential(AsyncContextManager):
     """
 
     def __init__(self, **kwargs: "Any") -> None:
+        super().__init__()
         self._refresh_token = None
         self._client = kwargs.pop("_client", None)
         self._tenant_id = kwargs.pop("tenant_id", None) or "organizations"
@@ -47,7 +48,6 @@ async def close(self):
         if self._client:
             await self._client.__aexit__()
 
-    @log_get_token_async
     async def get_token(self, *scopes, **kwargs):
         # type: (*str, **Any) -> AccessToken
         """Request an access token for `scopes` as the user currently signed in to Visual Studio Code.
@@ -59,29 +59,19 @@ async def get_token(self, *scopes, **kwargs):
         :raises ~azure.identity.CredentialUnavailableError: the credential cannot retrieve user details from Visual
           Studio Code
         """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
-
         if self._tenant_id.lower() == "adfs":
             raise CredentialUnavailableError(
                 message="VisualStudioCodeCredential authentication unavailable. ADFS is not supported."
             )
+        return await super().get_token(*scopes, **kwargs)
 
-        token = self._client.get_cached_access_token(scopes)
-        if not token:
-            token = await self._redeem_refresh_token(scopes, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                await self._redeem_refresh_token(scopes, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
+    async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]":
+        return self._client.get_cached_access_token(scopes)
 
-    async def _redeem_refresh_token(self, scopes: "Iterable[str]", **kwargs: "Any") -> "Optional[AccessToken]":
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         if not self._refresh_token:
             self._refresh_token = get_credentials()
             if not self._refresh_token:
                 raise CredentialUnavailableError(message="Failed to get Azure user details from Visual Studio Code.")
 
-        token = await self._client.obtain_token_by_refresh_token(scopes, self._refresh_token, **kwargs)
-        return token
+        return await self._client.obtain_token_by_refresh_token(scopes, self._refresh_token, **kwargs)

From eb164ea6bde6f340ba5ccf1b09c43b156e17760f Mon Sep 17 00:00:00 2001
From: Charles Lowell <chlowe@microsoft.com>
Date: Wed, 31 Mar 2021 14:49:26 -0700
Subject: [PATCH 2/2] update tests

---
 .../azure-identity/tests/test_aad_client.py   | 21 ----------------
 .../tests/test_aad_client_async.py            | 21 ----------------
 .../azure-identity/tests/test_authn_client.py | 21 ----------------
 .../tests/test_authn_client_async.py          | 21 ----------------
 .../tests/test_vscode_credential.py           | 22 +++++++++++------
 .../tests/test_vscode_credential_async.py     | 24 ++++++++++++-------
 6 files changed, 31 insertions(+), 99 deletions(-)

diff --git a/sdk/identity/azure-identity/tests/test_aad_client.py b/sdk/identity/azure-identity/tests/test_aad_client.py
index 88671fcc398b..4a313608a3b1 100644
--- a/sdk/identity/azure-identity/tests/test_aad_client.py
+++ b/sdk/identity/azure-identity/tests/test_aad_client.py
@@ -207,27 +207,6 @@ def send(request, **_):
     assert len(cache.find(TokenCache.CredentialType.REFRESH_TOKEN, query={"secret": invalid_token})) == 0
 
 
-def test_should_refresh():
-    client = AadClient("test", "test")
-    now = int(time.time())
-
-    # do not need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET + 1)
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
-
-    # need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    should_refresh = client.should_refresh(token)
-    assert should_refresh
-
-    # not exceed cool down time, do not refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    client._last_refresh_time = now - DEFAULT_TOKEN_REFRESH_RETRY_DELAY + 1
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
-
-
 def test_retries_token_requests():
     """The client should retry token requests"""
 
diff --git a/sdk/identity/azure-identity/tests/test_aad_client_async.py b/sdk/identity/azure-identity/tests/test_aad_client_async.py
index c85e0827623a..81d3fc2d827f 100644
--- a/sdk/identity/azure-identity/tests/test_aad_client_async.py
+++ b/sdk/identity/azure-identity/tests/test_aad_client_async.py
@@ -215,27 +215,6 @@ async def send(request, **_):
     assert len(cache.find(TokenCache.CredentialType.REFRESH_TOKEN, query={"secret": invalid_token})) == 0
 
 
-async def test_should_refresh():
-    client = AadClient("test", "test")
-    now = int(time.time())
-
-    # do not need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET + 1)
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
-
-    # need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    should_refresh = client.should_refresh(token)
-    assert should_refresh
-
-    # not exceed cool down time, do not refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    client._last_refresh_time = now - DEFAULT_TOKEN_REFRESH_RETRY_DELAY + 1
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
-
-
 async def test_retries_token_requests():
     """The client should retry token requests"""
 
diff --git a/sdk/identity/azure-identity/tests/test_authn_client.py b/sdk/identity/azure-identity/tests/test_authn_client.py
index 3dac71b924c7..ecab2c8a462d 100644
--- a/sdk/identity/azure-identity/tests/test_authn_client.py
+++ b/sdk/identity/azure-identity/tests/test_authn_client.py
@@ -233,24 +233,3 @@ def mock_send(request, **kwargs):
         client.request_token(("scope",))
         request = client.get_refresh_token_grant_request({"secret": "***"}, "scope")
         validate_url(request.url)
-
-
-def test_should_refresh():
-    client = AuthnClient(endpoint="http://foo")
-    now = int(time.time())
-
-    # do not need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET + 1)
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
-
-    # need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    should_refresh = client.should_refresh(token)
-    assert should_refresh
-
-    # not exceed cool down time, do not refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    client._last_refresh_time = now - DEFAULT_TOKEN_REFRESH_RETRY_DELAY + 1
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
diff --git a/sdk/identity/azure-identity/tests/test_authn_client_async.py b/sdk/identity/azure-identity/tests/test_authn_client_async.py
index 866f60c01ff4..6b6178163580 100644
--- a/sdk/identity/azure-identity/tests/test_authn_client_async.py
+++ b/sdk/identity/azure-identity/tests/test_authn_client_async.py
@@ -36,24 +36,3 @@ def mock_send(request, **kwargs):
     with patch.dict("os.environ", {EnvironmentVariables.AZURE_AUTHORITY_HOST: authority}, clear=True):
         client = AsyncAuthnClient(tenant=tenant_id, transport=Mock(send=wrap_in_future(mock_send)))
         await client.request_token(("scope",))
-
-
-def test_should_refresh():
-    client = AsyncAuthnClient(endpoint="http://foo")
-    now = int(time.time())
-
-    # do not need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET + 1)
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
-
-    # need refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    should_refresh = client.should_refresh(token)
-    assert should_refresh
-
-    # not exceed cool down time, do not refresh
-    token = AccessToken("token", now + DEFAULT_REFRESH_OFFSET - 1)
-    client._last_refresh_time = now - DEFAULT_TOKEN_REFRESH_RETRY_DELAY + 1
-    should_refresh = client.should_refresh(token)
-    assert not should_refresh
diff --git a/sdk/identity/azure-identity/tests/test_vscode_credential.py b/sdk/identity/azure-identity/tests/test_vscode_credential.py
index 2bd0a6ae7316..094deb76432d 100644
--- a/sdk/identity/azure-identity/tests/test_vscode_credential.py
+++ b/sdk/identity/azure-identity/tests/test_vscode_credential.py
@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import sys
+import time
 
 from azure.core.credentials import AccessToken
 from azure.identity import CredentialUnavailableError, VisualStudioCodeCredential
@@ -142,16 +143,23 @@ def test_cache_refresh_token():
 
 
 def test_no_obtain_token_if_cached():
-    expected_token = AccessToken("token", 42)
+    expected_token = AccessToken("token", time.time() + 3600)
 
-    mock_client = mock.Mock(should_refresh=lambda _: False)
-    mock_client.obtain_token_by_refresh_token = mock.Mock(return_value=expected_token)
-    mock_client.get_cached_access_token = mock.Mock(return_value="VALUE")
+    mock_client = mock.Mock(
+        obtain_token_by_refresh_token=mock.Mock(return_value=expected_token),
+        get_cached_access_token=mock.Mock(return_value=expected_token)
+    )
 
-    with mock.patch(VisualStudioCodeCredential.__module__ + ".get_credentials", return_value="VALUE"):
-        credential = VisualStudioCodeCredential(_client=mock_client)
+    credential = VisualStudioCodeCredential(_client=mock_client)
+    with mock.patch(
+        VisualStudioCodeCredential.__module__ + ".get_credentials",
+        mock.Mock(side_effect=Exception("credential should not acquire a new token")),
+    ):
         token = credential.get_token("scope")
-        assert mock_client.obtain_token_by_refresh_token.call_count == 0
+
+    assert mock_client.obtain_token_by_refresh_token.call_count == 0
+    assert token.token == expected_token.token
+    assert token.expires_on == expected_token.expires_on
 
 
 @pytest.mark.skipif(not sys.platform.startswith("linux"), reason="This test only runs on Linux")
diff --git a/sdk/identity/azure-identity/tests/test_vscode_credential_async.py b/sdk/identity/azure-identity/tests/test_vscode_credential_async.py
index f66b0318cbf3..898b14bd57d7 100644
--- a/sdk/identity/azure-identity/tests/test_vscode_credential_async.py
+++ b/sdk/identity/azure-identity/tests/test_vscode_credential_async.py
@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import time
 from unittest import mock
 from urllib.parse import urlparse
 
@@ -14,7 +15,7 @@
 import pytest
 
 from helpers import build_aad_response, mock_response, Request
-from helpers_async import async_validating_transport, AsyncMockTransport, wrap_in_future
+from helpers_async import async_validating_transport, wrap_in_future
 
 
 def test_tenant_id_validation():
@@ -145,17 +146,24 @@ async def test_cache_refresh_token():
 
 @pytest.mark.asyncio
 async def test_no_obtain_token_if_cached():
-    expected_token = AccessToken("token", 42)
+    expected_token = AccessToken("token", time.time() + 3600)
 
-    mock_client = mock.Mock(should_refresh=lambda _: False)
     token_by_refresh_token = mock.Mock(return_value=expected_token)
-    mock_client.obtain_token_by_refresh_token = wrap_in_future(token_by_refresh_token)
-    mock_client.get_cached_access_token = mock.Mock(return_value="VALUE")
+    mock_client = mock.Mock(
+        get_cached_access_token=mock.Mock(return_value=expected_token),
+        obtain_token_by_refresh_token=wrap_in_future(token_by_refresh_token)
+    )
 
-    with mock.patch(VisualStudioCodeCredential.__module__ + ".get_credentials", return_value="VALUE"):
-        credential = VisualStudioCodeCredential(_client=mock_client)
+    credential = VisualStudioCodeCredential(_client=mock_client)
+    with mock.patch(
+        VisualStudioCodeCredential.__module__ + ".get_credentials",
+        mock.Mock(side_effect=Exception("credential should not acquire a new token")),
+    ):
         token = await credential.get_token("scope")
-        assert token_by_refresh_token.call_count == 0
+
+    assert token_by_refresh_token.call_count == 0
+    assert token.token == expected_token.token
+    assert token.expires_on == expected_token.expires_on
 
 
 @pytest.mark.asyncio