support custom auth

Author: aviau

src/flareio/api_client.py

@@ -1,28 +1,30 @@
 import os
 
-from datetime import datetime
 from datetime import timedelta
 from http.cookiejar import DefaultCookiePolicy
 from urllib.parse import urljoin
 from urllib.parse import urlparse
 
 import requests
+from requests.auth import AuthBase
 
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
 
 import typing as t
 
-from flareio.exceptions import TokenError
+from flareio.auth import _FlareApiKeyAuth
 from flareio.models import ScrollEventsResult
 from flareio.ratelimit import Limiter
 from flareio.version import __version__ as _flareio_version
 
 
 _API_DOMAIN_DEFAULT: str = "api.flare.io"
-_ALLOWED_API_DOMAINS: t.Tuple[str, ...] = (
-    _API_DOMAIN_DEFAULT,
-    "api.eu.flare.io",
+_ALLOWED_API_DOMAINS: frozenset[str] = frozenset(
+    {
+        _API_DOMAIN_DEFAULT,
+        "api.eu.flare.io",
+    }
 )
 
 
@@ -34,7 +36,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: t.Optional[AuthBase] = None,
         _enable_beta_features: bool = False,
     ) -> None:
         if not api_key:
@@ -49,14 +51,16 @@ def __init__(
             raise Exception("Custom API domains considered a beta feature.")
         self._api_domain: str = api_domain
 
-        self._api_key: str = api_key
-        self._tenant_id: t.Optional[int] = tenant_id
-
-        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()
 
+        _auth = _auth or _FlareApiKeyAuth(
+            api_key=api_key,
+            api_domain=self._api_domain,
+            tenant_id=tenant_id,
+            session=self._session,
+        )
+        self._session.auth = _auth
+
     @classmethod
     def from_env(cls) -> "FlareApiClient":
         api_key: t.Optional[str] = os.environ.get("FLARE_API_KEY")
@@ -110,41 +114,7 @@ def _create_retry() -> Retry:
         return retry
 
     def generate_token(self) -> str:
-        payload: t.Optional[dict] = None
-
-        if self._tenant_id is not None:
-            payload = {
-                "tenant_id": self._tenant_id,
-            }
-
-        resp = self._session.post(
-            f"https://{self._api_domain}/tokens/generate",
-            json=payload,
-            headers={
-                "Authorization": self._api_key,
-            },
-        )
-        try:
-            resp.raise_for_status()
-        except Exception as ex:
-            raise TokenError("Failed to fetch API Token") from ex
-        token: str = resp.json()["token"]
-
-        self._api_token = token
-        self._api_token_exp = datetime.now() + timedelta(minutes=45)
-
-        return token
-
-    def _auth_headers(self) -> dict:
-        if self._disable_auth:
-            return dict()
-        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}"}
+        return self._auth.generate_token()
 
     def _request(
         self,
@@ -163,11 +133,6 @@ 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(
             method=method,
             url=url,

src/flareio/auth.py

@@ -0,0 +1,71 @@
+from datetime import datetime
+from datetime import timedelta
+
+import requests
+
+from requests.auth import AuthBase
+
+import typing as t
+
+from flareio.exceptions import TokenError
+
+
+class _FlareApiKeyAuth(AuthBase):
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        api_domain: str,
+        tenant_id: t.Optional[int] = None,
+        session: requests.Session,
+    ) -> None:
+        self._api_key: str = api_key
+        self._api_domain: str = api_domain
+        self._tenant_id: t.Optional[int] = tenant_id
+        self._session: requests.Session = session
+
+        self._api_token: t.Optional[str] = None
+        self._api_token_exp: t.Optional[datetime] = None
+
+    def generate_token(self) -> str:
+        payload: t.Optional[dict] = None
+
+        if self._tenant_id is not None:
+            payload = {
+                "tenant_id": self._tenant_id,
+            }
+
+        resp = self._session.post(
+            f"https://{self._api_domain}/tokens/generate",
+            json=payload,
+            headers={
+                "Authorization": self._api_key,
+            },
+        )
+        try:
+            resp.raise_for_status()
+        except Exception as ex:
+            raise TokenError("Failed to fetch API Token") from ex
+        token: str = resp.json()["token"]
+
+        self._api_token = token
+        self._api_token_exp = datetime.now() + timedelta(minutes=45)
+
+        return token
+
+    def __call__(
+        self,
+        r: requests.PreparedRequest,
+    ) -> requests.PreparedRequest:
+        # Token generation uses API key auth, don't override it.
+        if "Authorization" in r.headers:
+            return r
+
+        # Lazy token refresh.
+        if not self._api_token or (
+            self._api_token_exp and self._api_token_exp < datetime.now()
+        ):
+            self.generate_token()
+
+        r.headers["Authorization"] = f"Bearer {self._api_token}"
+        return r

tests/test_api_client_creation.py

@@ -38,8 +38,8 @@ def test_create_client_empty_api_key() -> None:
 
 def test_generate_token() -> None:
     client = get_test_client(authenticated=False)
-    assert client._api_token is None
-    assert client._api_token_exp is None
+    assert client._auth._api_token is None
+    assert client._auth._api_token_exp is None
     with requests_mock.Mocker() as mocker:
         mocker.register_uri(
             "POST",
@@ -53,9 +53,9 @@ def test_generate_token() -> None:
         token = client.generate_token()
         assert token == "test-token-hello"
 
-        assert client._api_token == "test-token-hello"
-        assert client._api_token_exp
-        assert client._api_token_exp >= datetime.now()
+        assert client._auth._api_token == "test-token-hello"
+        assert client._auth._api_token_exp
+        assert client._auth._api_token_exp >= datetime.now()
 
         assert mocker.last_request.url == "https://api.flare.io/tokens/generate"
         assert mocker.last_request.text is None
@@ -64,8 +64,8 @@ def test_generate_token() -> None:
 
 def test_generate_token_error() -> None:
     client = get_test_client(authenticated=False)
-    assert client._api_token is None
-    assert client._api_token_exp is None
+    assert client._auth._api_token is None
+    assert client._auth._api_token_exp is None
 
     with requests_mock.Mocker() as mocker:
         mocker.register_uri(

tests/test_api_client_endpoints.py

@@ -6,8 +6,8 @@
 
 def test_wrapped_methods() -> None:
     client = get_test_client(authenticated=False)
-    assert client._api_token is None
-    assert client._api_token_exp is None
+    assert client._auth._api_token is None
+    assert client._auth._api_token_exp is None
 
     # POST: This one will generate since its the first one.
     with requests_mock.Mocker() as mocker: