diff --git a/pyproject.toml b/pyproject.toml index affee54..e9457fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "python-squarelet" -version = "0.3.1" +version = "0.4.0" authors = [ { name="duckduckgrayduck", email="sanjin@muckrock.com" }, ] diff --git a/src/squarelet/squarelet.py b/src/squarelet/squarelet.py index 6d4bb54..370759a 100644 --- a/src/squarelet/squarelet.py +++ b/src/squarelet/squarelet.py @@ -47,9 +47,13 @@ def __init__( self.access_token = None self.refresh_token = None self._user_id = None - # Default UA for unauthenticated requests. - existing_ua = self.session.headers.get("User-Agent", "") - self.session.headers.update({"User-Agent": f"{existing_ua} Anonymous".strip()}) + self._username_cache = None + self._resolving_identity = False + + # Capture the library default UA once, before appending any identity, + # so the identity segment can be rebuilt deterministically later. + self._base_ua = self.session.headers.get("User-Agent", "") + self._set_user_agent() self._set_tokens() # Apply rate limiting @@ -63,6 +67,46 @@ def __init__( if rate_limit_sleep: self.request = ratelimit.sleep_and_retry(self.request) + def _fetch_me(self): + """Fetch the current user record from the API (used for identity/UA).""" + # set_tokens=False so a 403/429 here can't re-enter _set_tokens and loop. + return self.request("get", "users/me/", set_tokens=False).json() + + def _resolve_identity(self): + """Label for the UA: local username, cached value, else resolve from API. + + Token-authed clients (constructed without a username) have no local + identity, so the only source of truth is the token itself. We ask the + API who it belongs to and cache the result. + """ + if self.username: + return self.username + if self._username_cache is not None: + return self._username_cache + # Only hit the API if we're authenticated and not already resolving. + if not self.access_token or self._resolving_identity: + return None + self._resolving_identity = True + try: + data = self._fetch_me() + # Cache the user id too, so the user_id property doesn't re-fetch. + if self._user_id is None and data.get("id") is not None: + self._user_id = data["id"] + self._username_cache = data.get("username") or str(data.get("id")) + return self._username_cache + except Exception: # pylint: disable=broad-except + # UA labeling is cosmetic and must never break a real request. + return None + finally: + self._resolving_identity = False + + def _set_user_agent(self): + """Rebuild the User-Agent from the captured base plus current identity.""" + identity = self._resolve_identity() or "Anonymous" + self.session.headers.update( + {"User-Agent": f"{self._base_ua} {identity}".strip()} + ) + def _set_tokens(self): """Set the refresh and access tokens""" if self.refresh_token: @@ -76,15 +120,16 @@ def _set_tokens(self): else: self.access_token = None self.refresh_token = None + if self.access_token: self.session.headers.update( {"Authorization": f"Bearer {self.access_token}"} ) - # Identify authed users to better manage API usage. - if self.username: - existing_ua = self.session.headers.get("User-Agent", "") - new_ua = existing_ua.replace("Anonymous", self.username).strip() - self.session.headers.update({"User-Agent": new_ua}) + + # Rebuild the UA whenever auth state changes. This works for + # username/password clients AND token/refresh clients, since it can + # resolve identity from the API when no local username is available. + self._set_user_agent() def _get_tokens(self, username, password): """Get an access and refresh token in exchange for the username and password""" @@ -197,7 +242,6 @@ def raise_for_status(self, response): def user_id(self): """Returns the user ID of the user""" if self._user_id is None: - user_data = self.request("get", "users/me/").json() - user_id = user_data["id"] - self._user_id = user_id + user_data = self._fetch_me() + self._user_id = user_data["id"] return self._user_id diff --git a/src/squarelet/tests.py b/src/squarelet/tests.py index e1f4ffc..9a913b4 100644 --- a/src/squarelet/tests.py +++ b/src/squarelet/tests.py @@ -2,7 +2,6 @@ import os import time - import pytest from squarelet import CredentialsFailedError, DoesNotExistError, SquareletClient @@ -135,14 +134,88 @@ def test_user_agent_authenticated(squarelet_client): assert sq_user in ua assert "Anonymous" not in ua + def test_no_credentials_no_tokens(): """Test that a client without credentials has no tokens set""" client = SquareletClient(base_uri="https://api.www.documentcloud.org/api/") assert client.access_token is None assert client.refresh_token is None -def test_user_id_cached(squarelet_client): - """Test that user_id is fetched once and cached""" + +def test_user_id_fetched_only_once(squarelet_client): + """user_id should hit the API at most once, then serve from cache.""" + # Start from a clean slate so we're testing the fetch-and-cache path, + # not a value that construction may have already populated. + squarelet_client._user_id = None + + # Count how many times the API is actually hit. + call_count = {"n": 0} + real_fetch_me = squarelet_client._fetch_me + + def counting_fetch_me(): + call_count["n"] += 1 + return real_fetch_me() + + squarelet_client._fetch_me = counting_fetch_me + + # Read the property several times. first = squarelet_client.user_id - squarelet_client.session = None # Would blow up if a request was made - assert squarelet_client.user_id == first \ No newline at end of file + second = squarelet_client.user_id + third = squarelet_client.user_id + + # Same value every time, and the API was hit exactly once. + assert first == second == third + assert call_count["n"] == 1 + + +def test_user_agent_token_reuse_client(squarelet_client): + """A client authed via token reuse (no local username) should still be + labeled with the username, not 'Anonymous'. + + Regression: previously such clients stayed '... Anonymous' because the UA + relabel was gated on self.username being set. Mirrors the integration + pattern of reusing a refresh token on a freshly constructed client. + """ + sq_user = os.environ.get("SQ_USER") + + # A fresh client with no credentials starts as Anonymous. + rebuilt = SquareletClient(base_uri="https://api.www.documentcloud.org/api/") + assert "Anonymous" in rebuilt.session.headers["User-Agent"] + + # Reuse a valid refresh token from the authenticated fixture, then re-auth. + rebuilt.refresh_token = squarelet_client.refresh_token + rebuilt._set_tokens() + + ua = rebuilt.session.headers["User-Agent"] + assert "Anonymous" not in ua + assert sq_user in ua + # And it is genuinely authenticated. + assert rebuilt.session.headers.get("Authorization", "").startswith("Bearer ") + + +def test_user_agent_stable_across_reauth(squarelet_client): + """Repeated token refreshes must not corrupt or duplicate the UA label.""" + sq_user = os.environ.get("SQ_USER") + ua_before = squarelet_client.session.headers["User-Agent"] + squarelet_client._set_tokens() + squarelet_client._set_tokens() + ua_after = squarelet_client.session.headers["User-Agent"] + assert ua_before == ua_after + assert ua_after.count(sq_user) == 1 + + +def test_credentialed_client_does_not_call_api_for_label(squarelet_client): + """A client with a local username must label from it, not via users/me/.""" + sq_user = os.environ.get("SQ_USER") + calls = 0 + real = squarelet_client._fetch_me + + def counting(): + nonlocal calls + calls += 1 + return real() + + squarelet_client._fetch_me = counting + squarelet_client._set_tokens() # re-auth with creds present + assert sq_user in squarelet_client.session.headers["User-Agent"] + assert calls == 0