adds session to OAuth constructor so that we can mock the session to …

…run on CI.
dlt-hub · Jun 20, 2024 · edfff5b · edfff5b
1 parent 352b5d7
commit edfff5b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/dlt/sources/helpers/rest_client/auth.py b/dlt/sources/helpers/rest_client/auth.py
@@ -164,6 +164,7 @@ def __init__(
         client_secret: TSecretStrValue,
         access_token_request_data: Dict[str, Any] = None,
         default_token_expiration: int = 3600,
+        session: Annotated[BaseSession, NotResolved()] = None,
     ) -> None:
         super().__init__()
         self.access_token_url = access_token_url
@@ -175,6 +176,8 @@ def __init__(
             self.access_token_request_data = access_token_request_data
         self.default_token_expiration = default_token_expiration
         self.token_expiry: pendulum.DateTime = pendulum.now()
+        if session is None:
+            self.session = requests.client.session
 
     def __call__(self, request: PreparedRequest) -> PreparedRequest:
         if self.access_token is None or self.is_token_expired():
@@ -233,7 +236,7 @@ def __post_init__(self) -> None:
         self.scopes = self.scopes if isinstance(self.scopes, str) else " ".join(self.scopes)
         self.token = None
         self.token_expiry: Optional[pendulum.DateTime] = None
-        # use default system session is not specified
+        # use default system session unless specified otherwise
         if self.session is None:
             self.session = requests.client.session
 

diff --git a/tests/sources/helpers/rest_client/test_client.py b/tests/sources/helpers/rest_client/test_client.py
@@ -55,6 +55,7 @@ def rest_client_oauth() -> RESTClient:
         access_token_url=cast(TSecretStrValue, "https://api.example.com/oauth/token"),
         client_id=cast(TSecretStrValue, "test-client-id"),
         client_secret=cast(TSecretStrValue, "test-client-secret"),
+        session=Client().session,
     )
     return build_rest_client(auth=auth)
 
@@ -65,6 +66,7 @@ def rest_client_immediate_oauth_expiry(auth=None) -> RESTClient:
         access_token_url=cast(TSecretStrValue, "https://api.example.com/oauth/token-expires-now"),
         client_id=cast(TSecretStrValue, "test-client-id"),
         client_secret=cast(TSecretStrValue, "test-client-secret"),
+        session=Client().session,
     )
     return build_rest_client(auth=credentials_expiring_now)
 
@@ -207,6 +209,7 @@ def test_oauth2_client_credentials_flow_wrong_client_id(self, rest_client: RESTC
             access_token_url=cast(TSecretStrValue, "https://api.example.com/oauth/token"),
             client_id=cast(TSecretStrValue, "invalid-client-id"),
             client_secret=cast(TSecretStrValue, "test-client-secret"),
+            session=Client().session,
         )
 
         with pytest.raises(HTTPError) as e:
@@ -219,6 +222,7 @@ def test_oauth2_client_credentials_flow_wrong_client_secret(self, rest_client: R
             access_token_url=cast(TSecretStrValue, "https://api.example.com/oauth/token"),
             client_id=cast(TSecretStrValue, "test-client-id"),
             client_secret=cast(TSecretStrValue, "invalid-client-secret"),
+            session=Client().session,
         )
 
         with pytest.raises(HTTPError) as e:
@@ -274,6 +278,7 @@ def build_access_token_request(self) -> Dict[str, Any]:
             access_token_request_data={
                 "account_id": cast(TSecretStrValue, "test-account-id"),
             },
+            session=Client().session,
         )
 
         assert auth.build_access_token_request() == {