Add default_timeout for requests Session #510

lepture · Dec 6, 2022 · b0fc78f · b0fc78f
1 parent def72dc
commit b0fc78f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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