corehttp 1.0.0b5__py3-none-any.whl → 1.0.0b7__py3-none-any.whl

corehttp/rest/_rest_py3.py

@@ -300,7 +300,7 @@ class _HttpResponseBase(abc.ABC):
 
         :return: The JSON deserialized response body
         :rtype: any
-        :raises json.decoder.JSONDecodeError or ValueError (in python 2.7) if object is not JSON decodable:
+        :raises json.decoder.JSONDecodeError: if the body is not valid JSON.
         """
 
     @abc.abstractmethod
@@ -309,7 +309,7 @@ class _HttpResponseBase(abc.ABC):
 
         If response is good, does nothing.
 
-        :raises ~corehttp.HttpResponseError if the object has an error status code.:
+        :raises ~corehttp.HttpResponseError: if the object has an error status code.
         """
 
 
@@ -329,16 +329,13 @@ class HttpResponse(_HttpResponseBase):
     """
 
     @abc.abstractmethod
-    def __enter__(self) -> "HttpResponse":
-        ...
+    def __enter__(self) -> "HttpResponse": ...
 
     @abc.abstractmethod
-    def __exit__(self, *args: Any) -> None:
-        ...
+    def __exit__(self, *args: Any) -> None: ...
 
     @abc.abstractmethod
-    def close(self) -> None:
-        ...
+    def close(self) -> None: ...
 
     @abc.abstractmethod
     def read(self) -> bytes:
@@ -415,5 +412,4 @@ class AsyncHttpResponse(_HttpResponseBase, AsyncContextManager["AsyncHttpRespons
         yield  # pylint: disable=unreachable
 
     @abc.abstractmethod
-    async def close(self) -> None:
-        ...
+    async def close(self) -> None: ...

corehttp/runtime/pipeline/__init__.py

@@ -51,9 +51,7 @@ class PipelineContext(Dict[str, Any]):
 
     _PICKLE_CONTEXT = {"deserialized_data"}
 
-    def __init__(
-        self, transport: Optional["TransportType"], **kwargs: Any
-    ) -> None:  # pylint: disable=super-init-not-called
+    def __init__(self, transport: Optional["TransportType"], **kwargs: Any) -> None:
         self.transport: Optional["TransportType"] = transport
         self.options = kwargs
         self._protected = ["transport", "options"]
@@ -99,7 +97,7 @@ class PipelineContext(Dict[str, Any]):
     def clear(self) -> None:  # pylint: disable=docstring-missing-return, docstring-missing-rtype
         """Context objects cannot be cleared.
 
-        :raises: TypeError
+        :raises TypeError: If context objects cannot be cleared.
         """
         raise TypeError("Context objects cannot be cleared.")
 
@@ -108,17 +106,15 @@ class PipelineContext(Dict[str, Any]):
     ) -> None:
         """Context objects cannot be updated.
 
-        :raises: TypeError
+        :raises TypeError: If context objects cannot be updated.
         """
         raise TypeError("Context objects cannot be updated.")
 
     @overload
-    def pop(self, __key: str) -> Any:
-        ...
+    def pop(self, __key: str) -> Any: ...
 
     @overload
-    def pop(self, __key: str, __default: Optional[Any]) -> Any:
-        ...
+    def pop(self, __key: str, __default: Optional[Any]) -> Any: ...
 
     def pop(self, *args: Any) -> Any:
         """Removes specified key and returns the value.

corehttp/runtime/pipeline/_base.py

@@ -34,7 +34,7 @@ from . import (
     PipelineContext,
 )
 from ..policies import HTTPPolicy, SansIOHTTPPolicy
-from ._tools import await_result as _await_result
+from ._tools import sanitize_transport_options, await_result as _await_result
 from ...transport import HttpTransport
 
 HTTPResponseType = TypeVar("HTTPResponseType")
@@ -103,6 +103,7 @@ class _TransportRunner(HTTPPolicy[HTTPRequestType, HTTPResponseType]):
         :return: The PipelineResponse object.
         :rtype: ~corehttp.runtime.pipeline.PipelineResponse
         """
+        sanitize_transport_options(request.context.options)
         return PipelineResponse(
             request.http_request,
             self._sender.send(request.http_request, **request.context.options),
@@ -153,7 +154,7 @@ class Pipeline(ContextManager["Pipeline"], Generic[HTTPRequestType, HTTPResponse
         self._transport.__enter__()
         return self
 
-    def __exit__(self, *exc_details: Any) -> None:  # pylint: disable=arguments-differ
+    def __exit__(self, *exc_details: Any) -> None:
         self._transport.__exit__(*exc_details)
 
     def run(self, request: HTTPRequestType, **kwargs: Any) -> PipelineResponse[HTTPRequestType, HTTPResponseType]:

corehttp/runtime/pipeline/_base_async.py

@@ -26,13 +26,14 @@
 from __future__ import annotations
 import inspect
 from types import TracebackType
-from typing import Any, Union, Generic, TypeVar, List, Optional, Iterable, Type
-from typing_extensions import AsyncContextManager, TypeGuard
+from typing import Any, Union, Generic, TypeVar, List, Optional, Iterable, Type, AsyncContextManager
+from typing_extensions import TypeGuard
 
 from . import PipelineRequest, PipelineResponse, PipelineContext
 from ..policies import AsyncHTTPPolicy, SansIOHTTPPolicy
 from ..pipeline._base import is_sansio_http_policy
 from ._tools_async import await_result as _await_result
+from ._tools import sanitize_transport_options
 from ...transport import AsyncHttpTransport
 
 AsyncHTTPResponseType = TypeVar("AsyncHTTPResponseType")
@@ -45,9 +46,7 @@ def is_async_http_policy(policy: object) -> TypeGuard[AsyncHTTPPolicy]:
     return False
 
 
-class _SansIOAsyncHTTPPolicyRunner(
-    AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]
-):  # pylint: disable=unsubscriptable-object
+class _SansIOAsyncHTTPPolicyRunner(AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]):
     """Async implementation of the SansIO policy.
 
     Modifies the request and sends to the next policy in the chain.
@@ -76,9 +75,7 @@ class _SansIOAsyncHTTPPolicyRunner(
         return response
 
 
-class _AsyncTransportRunner(
-    AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]
-):  # pylint: disable=unsubscriptable-object
+class _AsyncTransportRunner(AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]):
     """Async Transport runner.
 
     Uses specified HTTP transport type to send request and returns response.
@@ -101,6 +98,7 @@ class _AsyncTransportRunner(
         :return: The PipelineResponse object.
         :rtype: ~corehttp.runtime.pipeline.PipelineResponse
         """
+        sanitize_transport_options(request.context.options)
         return PipelineResponse(
             request.http_request,
             await self._sender.send(request.http_request, **request.context.options),

corehttp/runtime/pipeline/_tools.py

@@ -23,7 +23,7 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-from typing import Callable, TypeVar
+from typing import Callable, TypeVar, Dict
 from typing_extensions import ParamSpec
 
 P = ParamSpec("P")
@@ -39,9 +39,25 @@ def await_result(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
     :type args: list
     :rtype: any
     :return: The result of the function
-    :raises: TypeError
+    :raises TypeError: If the function returns an awaitable object.
     """
     result = func(*args, **kwargs)
     if hasattr(result, "__await__"):
         raise TypeError("Policy {} returned awaitable object in non-async pipeline.".format(func))
     return result
+
+
+def sanitize_transport_options(options: Dict[str, str]) -> None:
+    """Remove options that could potentially make it to the transport layer.
+
+    - "tracing_options" is used in the DistributedHttpTracingPolicy and tracing decorators
+
+    :param options: The options.
+    :type options: dict
+    """
+    if not options:
+        return
+
+    options_to_remove = ["tracing_options"]
+    for key in options_to_remove:
+        options.pop(key, None)

corehttp/runtime/pipeline/_tools_async.py

@@ -31,13 +31,11 @@ T = TypeVar("T")
 
 
 @overload
-async def await_result(func: Callable[P, Awaitable[T]], *args: P.args, **kwargs: P.kwargs) -> T:
-    ...
+async def await_result(func: Callable[P, Awaitable[T]], *args: P.args, **kwargs: P.kwargs) -> T: ...
 
 
 @overload
-async def await_result(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
-    ...
+async def await_result(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: ...
 
 
 async def await_result(func: Callable[P, Union[T, Awaitable[T]]], *args: P.args, **kwargs: P.kwargs) -> T:

corehttp/runtime/policies/__init__.py

@@ -29,6 +29,7 @@ from ._authentication import (
     BearerTokenCredentialPolicy,
     ServiceKeyCredentialPolicy,
 )
+from ._distributed_tracing import DistributedHttpTracingPolicy
 from ._retry import RetryPolicy, RetryMode
 from ._universal import (
     HeadersPolicy,
@@ -57,4 +58,5 @@ __all__ = [
     "AsyncHTTPPolicy",
     "AsyncBearerTokenCredentialPolicy",
     "AsyncRetryPolicy",
+    "DistributedHttpTracingPolicy",
 ]

corehttp/runtime/policies/_authentication.py

@@ -5,16 +5,17 @@
 # -------------------------------------------------------------------------
 from __future__ import annotations
 import time
-from typing import TYPE_CHECKING, Optional, TypeVar, MutableMapping, Any
+from typing import TYPE_CHECKING, Optional, TypeVar, MutableMapping, Any, Union
 
+from ...credentials import TokenRequestOptions
 from ...rest import HttpResponse, HttpRequest
 from . import HTTPPolicy, SansIOHTTPPolicy
-from ...exceptions import ServiceRequestError
+from ...exceptions import ServiceRequestError, HttpResponseError
 
 if TYPE_CHECKING:
-    # pylint:disable=unused-import
+
     from ...credentials import (
-        AccessToken,
+        AccessTokenInfo,
         TokenCredential,
         ServiceKeyCredential,
     )
@@ -31,15 +32,23 @@ class _BearerTokenCredentialPolicyBase:
     :param credential: The credential.
     :type credential: ~corehttp.credentials.TokenCredential
     :param str scopes: Lets you specify the type of access needed.
+    :keyword auth_flows: A list of authentication flows to use for the credential.
+    :paramtype auth_flows: list[dict[str, Union[str, list[dict[str, str]]]]]
     """
 
+    # pylint: disable=unused-argument
     def __init__(
-        self, credential: "TokenCredential", *scopes: str, **kwargs: Any  # pylint: disable=unused-argument
+        self,
+        credential: "TokenCredential",
+        *scopes: str,
+        auth_flows: Optional[list[dict[str, Union[str, list[dict[str, str]]]]]] = None,
+        **kwargs: Any,
     ) -> None:
         super(_BearerTokenCredentialPolicyBase, self).__init__()
         self._scopes = scopes
         self._credential = credential
-        self._token: Optional["AccessToken"] = None
+        self._token: Optional["AccessTokenInfo"] = None
+        self._auth_flows = auth_flows
 
     @staticmethod
     def _enforce_https(request: PipelineRequest[HTTPRequestType]) -> None:
@@ -68,7 +77,12 @@ class _BearerTokenCredentialPolicyBase:
 
     @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]):
@@ -77,20 +91,33 @@ class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, HTTPPolicy[H
     :param credential: The credential.
     :type credential: ~corehttp.TokenCredential
     :param str scopes: Lets you specify the type of access needed.
-    :raises: :class:`~corehttp.exceptions.ServiceRequestError`
+    :keyword auth_flows: A list of authentication flows to use for the credential.
+    :paramtype auth_flows: list[dict[str, Union[str, list[dict[str, str]]]]]
+    :raises ~corehttp.exceptions.ServiceRequestError: If the request fails.
     """
 
-    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
+    def on_request(
+        self,
+        request: PipelineRequest[HTTPRequestType],
+        *,
+        auth_flows: Optional[list[dict[str, Union[str, list[dict[str, str]]]]]] = None,
+    ) -> None:
         """Called before the policy sends a request.
 
         The base implementation authorizes the request with a bearer token.
 
         :param ~corehttp.runtime.pipeline.PipelineRequest request: the request
+        :keyword auth_flows: A list of authentication flows to use for the credential.
+        :paramtype auth_flows: list[dict[str, Union[str, list[dict[str, str]]]]]
         """
+        # If auth_flows is an empty list, we should not attempt to authorize the request.
+        if auth_flows is not None and len(auth_flows) == 0:
+            return
         self._enforce_https(request)
 
         if self._token is None or self._need_new_token:
-            self._token = self._credential.get_token(*self._scopes)
+            options: TokenRequestOptions = {"auth_flows": auth_flows} if auth_flows else {}  # type: ignore
+            self._token = self._credential.get_token_info(*self._scopes, options=options)
         self._update_headers(request.http_request.headers, self._token.token)
 
     def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs: Any) -> None:
@@ -102,7 +129,12 @@ class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, HTTPPolicy[H
         :param ~corehttp.runtime.pipeline.PipelineRequest request: the request
         :param str scopes: required scopes of authentication
         """
-        self._token = self._credential.get_token(*scopes, **kwargs)
+        options: TokenRequestOptions = {}
+        # Loop through all the keyword arguments and check if they are part of the TokenRequestOptions.
+        for key in list(kwargs.keys()):
+            if key in TokenRequestOptions.__annotations__:  # pylint:disable=no-member
+                options[key] = kwargs.pop(key)  # type: ignore[literal-required]
+        self._token = self._credential.get_token_info(*scopes, options=options)
         self._update_headers(request.http_request.headers, self._token.token)
 
     def send(self, request: PipelineRequest[HTTPRequestType]) -> PipelineResponse[HTTPRequestType, HTTPResponseType]:
@@ -113,25 +145,39 @@ class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, HTTPPolicy[H
         :return: The pipeline response object
         :rtype: ~corehttp.runtime.pipeline.PipelineResponse
         """
-        self.on_request(request)
+        op_auth_flows = request.context.options.pop("auth_flows", None)
+        auth_flows = op_auth_flows if op_auth_flows is not None else self._auth_flows
+        self.on_request(request, auth_flows=auth_flows)
         try:
             response = self.next.send(request)
-            self.on_response(request, response)
-        except Exception:  # pylint:disable=broad-except
+        except Exception:
             self.on_exception(request)
             raise
-        else:
-            if response.http_response.status_code == 401:
-                self._token = None  # any cached token is invalid
-                if "WWW-Authenticate" in response.http_response.headers:
+
+        self.on_response(request, response)
+        if response.http_response.status_code == 401:
+            self._token = None  # any cached token is invalid
+            if "WWW-Authenticate" in response.http_response.headers:
+                try:
                     request_authorized = self.on_challenge(request, response)
-                    if request_authorized:
+                except Exception as ex:
+                    # If the response is streamed, read it so the error message is immediately available to the user.
+                    # Otherwise, a generic error message will be given and the user will have to read the response
+                    # body to see the actual error.
+                    if response.context.options.get("stream"):
                         try:
-                            response = self.next.send(request)
-                            self.on_response(request, response)
+                            response.http_response.read()  # type: ignore
                         except Exception:  # pylint:disable=broad-except
-                            self.on_exception(request)
-                            raise
+                            pass
+                    # Raise the exception from the token request with the original 401 response.
+                    raise ex from HttpResponseError(response=response.http_response)
+                if request_authorized:
+                    try:
+                        response = self.next.send(request)
+                        self.on_response(request, response)
+                    except Exception:
+                        self.on_exception(request)
+                        raise
 
         return response
 
@@ -181,7 +227,8 @@ class ServiceKeyCredentialPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseT
     :type credential: ~corehttp.credentials.ServiceKeyCredential
     :param str name: The name of the key header used for the credential.
     :keyword str prefix: The name of the prefix for the header value if any.
-    :raises: ValueError or TypeError
+    :raises ValueError: if name is None or empty.
+    :raises TypeError: if name is not a string or if credential is not an instance of ServiceKeyCredential.
     """
 
     def __init__(  # pylint: disable=unused-argument
@@ -204,4 +251,9 @@ class ServiceKeyCredentialPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseT
         self._prefix = prefix + " " if prefix else ""
 
     def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
+        """Called before the policy sends a request.
+
+        :param request: The request to be modified before sending.
+        :type request: ~corehttp.runtime.pipeline.PipelineRequest
+        """
         request.http_request.headers[self._name] = f"{self._prefix}{self._credential.key}"

corehttp/runtime/policies/_authentication_async.py

@@ -5,14 +5,15 @@
 # -------------------------------------------------------------------------
 from __future__ import annotations
 import time
-from typing import TYPE_CHECKING, Any, Awaitable, Optional, cast, TypeVar
+from typing import TYPE_CHECKING, Any, Awaitable, Optional, cast, TypeVar, Union
 
-from ...credentials import AccessToken
+from ...credentials import AccessTokenInfo, TokenRequestOptions
 from ..pipeline import PipelineRequest, PipelineResponse
 from ..pipeline._tools_async import await_result
 from ._base_async import AsyncHTTPPolicy
 from ._authentication import _BearerTokenCredentialPolicyBase
 from ...rest import AsyncHttpResponse, HttpRequest
+from ...exceptions import HttpResponseError
 from ...utils._utils import get_running_async_lock
 
 if TYPE_CHECKING:
@@ -29,16 +30,24 @@ class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy[HTTPRequestType, AsyncHTT
     :param credential: The credential.
     :type credential: ~corehttp.credentials.TokenCredential
     :param str scopes: Lets you specify the type of access needed.
+    :keyword auth_flows: A list of authentication flows to use for the credential.
+    :paramtype auth_flows: list[dict[str, Union[str, list[dict[str, str]]]]]
     """
 
+    # pylint: disable=unused-argument
     def __init__(
-        self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: Any  # pylint: disable=unused-argument
+        self,
+        credential: "AsyncTokenCredential",
+        *scopes: str,
+        auth_flows: Optional[list[dict[str, Union[str, list[dict[str, str]]]]]] = None,
+        **kwargs: Any,
     ) -> None:
         super().__init__()
         self._credential = credential
         self._lock_instance = None
         self._scopes = scopes
-        self._token: Optional["AccessToken"] = None
+        self._token: Optional[AccessTokenInfo] = None
+        self._auth_flows = auth_flows
 
     @property
     def _lock(self):
@@ -46,21 +55,32 @@ class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy[HTTPRequestType, AsyncHTT
             self._lock_instance = get_running_async_lock()
         return self._lock_instance
 
-    async def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
+    async def on_request(
+        self,
+        request: PipelineRequest[HTTPRequestType],
+        *,
+        auth_flows: Optional[list[dict[str, Union[str, list[dict[str, str]]]]]] = None,
+    ) -> None:
         """Adds a bearer token Authorization header to request and sends request to next policy.
 
         :param request: The pipeline request object to be modified.
         :type request: ~corehttp.runtime.pipeline.PipelineRequest
-        :raises: :class:`~corehttp.exceptions.ServiceRequestError`
+        :keyword auth_flows: A list of authentication flows to use for the credential.
+        :paramtype auth_flows: list[dict[str, Union[str, list[dict[str, str]]]]]
+        :raises ~corehttp.exceptions.ServiceRequestError: If the request fails.
         """
+        # If auth_flows is an empty list, we should not attempt to authorize the request.
+        if auth_flows is not None and len(auth_flows) == 0:
+            return
         _BearerTokenCredentialPolicyBase._enforce_https(request)  # pylint:disable=protected-access
 
-        if self._token is None or self._need_new_token():
+        if self._token is None or self._need_new_token:
             async with self._lock:
                 # double check because another coroutine may have acquired a token while we waited to acquire the lock
-                if self._token is None or self._need_new_token():
-                    self._token = await await_result(self._credential.get_token, *self._scopes)
-        request.http_request.headers["Authorization"] = "Bearer " + cast(AccessToken, self._token).token
+                if self._token is None or self._need_new_token:
+                    options: TokenRequestOptions = {"auth_flows": auth_flows} if auth_flows else {}  # type: ignore
+                    self._token = await await_result(self._credential.get_token_info, *self._scopes, options=options)
+        request.http_request.headers["Authorization"] = "Bearer " + cast(AccessTokenInfo, self._token).token
 
     async def authorize_request(self, request: PipelineRequest[HTTPRequestType], *scopes: str, **kwargs: Any) -> None:
         """Acquire a token from the credential and authorize the request with it.
@@ -71,9 +91,15 @@ class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy[HTTPRequestType, AsyncHTT
         :param ~corehttp.runtime.pipeline.PipelineRequest request: the request
         :param str scopes: required scopes of authentication
         """
+        options: TokenRequestOptions = {}
+        # Loop through all the keyword arguments and check if they are part of the TokenRequestOptions.
+        for key in list(kwargs.keys()):
+            if key in TokenRequestOptions.__annotations__:  # pylint:disable=no-member
+                options[key] = kwargs.pop(key)  # type: ignore[literal-required]
+
         async with self._lock:
-            self._token = await await_result(self._credential.get_token, *scopes, **kwargs)
-        request.http_request.headers["Authorization"] = "Bearer " + cast(AccessToken, self._token).token
+            self._token = await await_result(self._credential.get_token_info, *scopes, options=options)
+        request.http_request.headers["Authorization"] = "Bearer " + cast(AccessTokenInfo, self._token).token
 
     async def send(
         self, request: PipelineRequest[HTTPRequestType]
@@ -85,27 +111,40 @@ class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy[HTTPRequestType, AsyncHTT
         :return: The pipeline response object
         :rtype: ~corehttp.runtime.pipeline.PipelineResponse
         """
-        await await_result(self.on_request, request)
+        op_auth_flows = request.context.options.pop("auth_flows", None)
+        auth_flows = op_auth_flows if op_auth_flows is not None else self._auth_flows
+        await await_result(self.on_request, request, auth_flows=auth_flows)
         try:
             response = await self.next.send(request)
-        except Exception:  # pylint:disable=broad-except
+        except Exception:
             await await_result(self.on_exception, request)
             raise
-        else:
-            await await_result(self.on_response, request, response)
+        await await_result(self.on_response, request, response)
 
         if response.http_response.status_code == 401:
             self._token = None  # any cached token is invalid
             if "WWW-Authenticate" in response.http_response.headers:
-                request_authorized = await self.on_challenge(request, response)
+                try:
+                    request_authorized = await self.on_challenge(request, response)
+                except Exception as ex:
+                    # If the response is streamed, read it so the error message is immediately available to the user.
+                    # Otherwise, a generic error message will be given and the user will have to read the response
+                    # body to see the actual error.
+                    if response.context.options.get("stream"):
+                        try:
+                            await response.http_response.read()  # type: ignore
+                        except Exception:  # pylint:disable=broad-except
+                            pass
+
+                    # Raise the exception from the token request with the original 401 response
+                    raise ex from HttpResponseError(response=response.http_response)
                 if request_authorized:
                     try:
                         response = await self.next.send(request)
-                    except Exception:  # pylint:disable=broad-except
+                    except Exception:
                         await await_result(self.on_exception, request)
                         raise
-                    else:
-                        await await_result(self.on_response, request, response)
+                    await await_result(self.on_response, request, response)
 
         return response
 
@@ -151,5 +190,11 @@ class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy[HTTPRequestType, AsyncHTT
         # pylint: disable=unused-argument
         return
 
+    @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)
+        )