diff --git a/authlib/integrations/requests_client/assertion_session.py b/authlib/integrations/requests_client/assertion_session.py index b5eb3891..5d4e6bc7 100644 --- a/authlib/integrations/requests_client/assertion_session.py +++ b/authlib/integrations/requests_client/assertion_session.py @@ -25,8 +25,9 @@ class AssertionSession(AssertionClient, Session): DEFAULT_GRANT_TYPE = JWT_BEARER_GRANT_TYPE def __init__(self, token_endpoint, issuer, subject, audience=None, grant_type=None, - claims=None, token_placement='header', scope=None, **kwargs): + claims=None, token_placement='header', scope=None, default_timeout=None, **kwargs): Session.__init__(self) + self.default_timeout = default_timeout update_session_configure(self, kwargs) AssertionClient.__init__( self, session=self, @@ -37,6 +38,8 @@ def __init__(self, token_endpoint, issuer, subject, audience=None, grant_type=No def request(self, method, url, withhold_token=False, auth=None, **kwargs): """Send request with auto refresh token feature.""" + if self.default_timeout: + kwargs.setdefault('timeout', self.default_timeout) if not withhold_token and auth is None: auth = self.token_auth return super(AssertionSession, self).request( diff --git a/authlib/integrations/requests_client/oauth2_session.py b/authlib/integrations/requests_client/oauth2_session.py index 65db8427..3b468197 100644 --- a/authlib/integrations/requests_client/oauth2_session.py +++ b/authlib/integrations/requests_client/oauth2_session.py @@ -64,6 +64,7 @@ class OAuth2Session(OAuth2Client, Session): values: "header", "body", "uri". :param update_token: A function for you to update token. It accept a :class:`OAuth2Token` as parameter. + :param default_timeout: If settled, every requests will have a default timeout. """ client_auth_class = OAuth2ClientAuth token_auth_class = OAuth2Auth @@ -78,10 +79,9 @@ def __init__(self, client_id=None, client_secret=None, revocation_endpoint_auth_method=None, scope=None, state=None, redirect_uri=None, token=None, token_placement='header', - update_token=None, **kwargs): - - self.default_timeout = kwargs.get('timeout') + update_token=None, default_timeout=None, **kwargs): Session.__init__(self) + self.default_timeout = default_timeout update_session_configure(self, kwargs) OAuth2Client.__init__( @@ -100,7 +100,8 @@ def fetch_access_token(self, url=None, **kwargs): def request(self, method, url, withhold_token=False, auth=None, **kwargs): """Send request with auto refresh token feature (if available).""" - kwargs['timeout'] = kwargs.get('timeout') or self.default_timeout + if self.default_timeout: + kwargs.setdefault('timeout', self.default_timeout) if not withhold_token and auth is None: if not self.token: raise MissingTokenError() diff --git a/tests/clients/test_requests/test_oauth2_session.py b/tests/clients/test_requests/test_oauth2_session.py index 4c918f36..fd26da64 100644 --- a/tests/clients/test_requests/test_oauth2_session.py +++ b/tests/clients/test_requests/test_oauth2_session.py @@ -508,7 +508,7 @@ def verifier(r, **kwargs): sess.get('https://i.b', auth=client.token_auth) def test_use_default_request_timeout(self): - expected_timeout = 10 + expected_timeout = 15 def verifier(r, **kwargs): timeout = kwargs.get('timeout') @@ -519,7 +519,7 @@ def verifier(r, **kwargs): client = OAuth2Session( client_id=self.client_id, token=self.token, - timeout=expected_timeout, + default_timeout=expected_timeout, ) client.send = verifier @@ -538,7 +538,7 @@ def verifier(r, **kwargs): client = OAuth2Session( client_id=self.client_id, token=self.token, - timeout=default_timeout, + default_timeout=default_timeout, ) client.send = verifier