Resolve identity more robustly

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" },
 ]

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

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
+    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