add customizable auth

Author: aviau

src/flareio/api_client.py

@@ -9,6 +9,7 @@
 import requests
 
 from requests.adapters import HTTPAdapter
+from requests.auth import AuthBase
 from urllib3.util import Retry
 
 import typing as t
@@ -34,7 +35,7 @@ def __init__(
         tenant_id: t.Optional[int] = None,
         session: t.Optional[requests.Session] = None,
         api_domain: t.Optional[str] = None,
-        _disable_auth: bool = False,
+        _auth: AuthBase | None = None,
         _enable_beta_features: bool = False,
     ) -> None:
         if not api_key:
@@ -52,9 +53,9 @@ def __init__(
         self._api_key: str = api_key
         self._tenant_id: t.Optional[int] = tenant_id
 
+        self._auth: t.Optional[AuthBase] = _auth
         self._api_token: t.Optional[str] = None
         self._api_token_exp: t.Optional[datetime] = None
-        self._disable_auth: bool = _disable_auth
         self._session = session or self._create_session()
 
     @classmethod
@@ -135,16 +136,24 @@ def generate_token(self) -> str:
 
         return token
 
-    def _auth_headers(self) -> dict:
-        if self._disable_auth:
-            return dict()
+    def _apply_auth(
+        self,
+        *,
+        request: requests.PreparedRequest,
+    ) -> requests.PreparedRequest:
+        if self._auth:
+            self._auth(request)
+            return request
+
         api_token: t.Optional[str] = self._api_token
         if not api_token or (
             self._api_token_exp and self._api_token_exp < datetime.now()
         ):
             api_token = self.generate_token()
 
-        return {"Authorization": f"Bearer {api_token}"}
+        request.headers["Authorization"] = f"Bearer {api_token}"
+
+        return request
 
     def _request(
         self,
@@ -163,19 +172,20 @@ def _request(
                 f"Client was used to access {netloc=} at {url=}. Only the domain {self._api_domain} is supported."
             )
 
-        headers = {
-            **(headers or {}),
-            **self._auth_headers(),
-        }
-
-        return self._session.request(
+        request = requests.Request(
             method=method,
             url=url,
             params=params,
             json=json,
             headers=headers,
         )
 
+        prepared = self._session.prepare_request(request)
+        prepared = self._apply_auth(request=prepared)
+        resp = self._session.send(prepared)
+
+        return resp
+
     def post(
         self,
         url: str,

src/flareio/auth.py

@@ -0,0 +1,26 @@
+from requests import PreparedRequest
+from requests.auth import AuthBase
+
+
+class _StaticHeadersAuth(AuthBase):
+    def __init__(
+        self,
+        *,
+        headers: dict[str, str],
+    ) -> None:
+        self._headers: dict[str, str] = headers
+
+    def __call__(
+        self,
+        r: PreparedRequest,
+    ) -> PreparedRequest:
+        r.headers.update(self._headers)
+        return r
+
+
+class _EmptyAuth(AuthBase):
+    def __call__(
+        self,
+        r: PreparedRequest,
+    ) -> PreparedRequest:
+        return r

tests/test_api_client_endpoints.py

@@ -3,6 +3,8 @@
 
 from .utils import get_test_client
 
+from flareio.auth import _EmptyAuth
+
 
 def test_wrapped_methods() -> None:
     client = get_test_client(authenticated=False)
@@ -126,7 +128,7 @@ def test_bad_domain() -> None:
 def test_disable_auth_does_not_call_generate() -> None:
     client = get_test_client(
         authenticated=False,
-        _disable_auth=True,
+        _auth=_EmptyAuth(),
     )
     with requests_mock.Mocker() as mocker:
         mocker.register_uri(

tests/utils.py

@@ -1,5 +1,7 @@
 import requests_mock
 
+from requests.auth import AuthBase
+
 import typing as t
 
 from flareio import FlareApiClient
@@ -11,14 +13,14 @@ def get_test_client(
     authenticated: bool = True,
     api_domain: t.Optional[str] = None,
     _enable_beta_features: bool = False,
-    _disable_auth: bool = False,
+    _auth: t.Optional[AuthBase] = None,
 ) -> FlareApiClient:
     client = FlareApiClient(
         api_key="test-api-key",
         tenant_id=tenant_id,
         api_domain=api_domain,
         _enable_beta_features=_enable_beta_features,
-        _disable_auth=_disable_auth,
+        _auth=_auth,
     )
 
     if authenticated: