test_gemini_errors.py

"""Unit tests (no DB, no network) for how Gemini errors are surfaced (28 Sep 2026):

  * a 429 that means "budget spent" (RESOURCE_EXHAUSTED / monthly spending
    cap) fails at once — no 2/4/8 s retries — as GeminiHTTPError with
    .quota == True;
  * no error message ever carries the request URL's `?key=` (requests'
    HTTPError and urllib3's connection errors both quote the keyed URL,
    which is how the live API key reached the public event stream);
  * an ordinary 429 / 5xx still retries and then fails without the key;
  * batch_ingest turns the exception into a `gemini_error:` result with
    quota set, key-free;
  * agdb.redact() and its SQL twin scrub key / api_key / token parameters
    and bearer tokens, and leave already-redacted text alone.

  python test_gemini_errors.py
"""
import json
import os
import sys

os.environ["AGENT_USE_NTRS"] = "0"

import requests

import extraction
from agent import agdb

fails = []


def check(name, cond):
    print(("PASS " if cond else "FAIL ") + name)
    if not cond:
        fails.append(name)


KEYED = "https://generativelanguage.googleapis.com/v1beta/models/x:generateContent?key=AQ.SECRETSECRETSECRET"


class FakeResp:
    def __init__(self, status, body=None, headers=None, reason="Too Many Requests"):
        self.status_code = status
        self._body = body
        self.headers = headers or {}
        self.reason = reason
        self.text = json.dumps(body) if body is not None else ""

    def json(self):
        if self._body is None:
            raise ValueError("no json")
        return self._body


CALLS: list = []
SLEEPS: list = []
SCRIPT: list = []


def fake_request(method, url, timeout=None, **kw):
    CALLS.append(url)
    nxt = SCRIPT.pop(0) if SCRIPT else FakeResp(200, {})
    if isinstance(nxt, Exception):
        raise nxt
    return nxt


extraction.requests.request = fake_request
_sleep = lambda s: SLEEPS.append(s)

# --- 1. spending cap: fail fast, quota flag, no key --------------------------------
cap_body = {"error": {"code": 429, "status": "RESOURCE_EXHAUSTED",
                      "message": "Your project has exceeded its monthly spending cap. Please go to AI Studio at https://ai.studio/spend to manage your project spend cap."}}
SCRIPT[:] = [FakeResp(429, cap_body)]
CALLS.clear(); SLEEPS.clear()
try:
    extraction._request_with_retry("POST", KEYED, json={}, _sleep=_sleep)
    check("spending cap raises", False)
except extraction.GeminiHTTPError as exc:
    check("spending cap raises GeminiHTTPError at once (1 call, no sleep)", len(CALLS) == 1 and SLEEPS == [])
    check("…with quota == True and the API's status", exc.quota and exc.status == "RESOURCE_EXHAUSTED")
    check("…message carries the API's detail, not the keyed URL",
          "spending cap" in str(exc) and "key=" not in str(exc) and "SECRET" not in str(exc))
    check("…is still a requests.RequestException (batch_ingest's except clause)",
          isinstance(exc, requests.RequestException))

# --- 2. ordinary 429 (rate limit) still retries, then fails without the key ----------
SCRIPT[:] = [FakeResp(429, {"error": {"code": 429, "status": "UNAVAILABLE", "message": "slow down"}})] * 4
CALLS.clear(); SLEEPS.clear()
try:
    extraction._request_with_retry("POST", KEYED, json={}, _sleep=_sleep)
    check("plain 429 raises after retries", False)
except extraction.GeminiHTTPError as exc:
    check("plain 429: retried MAX_RETRIES times with backoff", len(CALLS) == extraction.MAX_RETRIES + 1 and len(SLEEPS) == extraction.MAX_RETRIES)
    check("plain 429: not a quota refusal", not exc.quota)
    check("plain 429: message key-free", "key=" not in str(exc) and "SECRET" not in str(exc))

# --- 3. 5xx then 200 recovers; 400 fails once, key-free ------------------------------
SCRIPT[:] = [FakeResp(503, None, reason="Service Unavailable"), FakeResp(200, {"ok": 1})]
CALLS.clear(); SLEEPS.clear()
r = extraction._request_with_retry("POST", KEYED, json={}, _sleep=_sleep)
check("503 then 200: returns the 200 after one retry", r is not None and r.status_code == 200 and len(CALLS) == 2)
SCRIPT[:] = [FakeResp(400, {"error": {"code": 400, "status": "INVALID_ARGUMENT", "message": "bad schema"}}, reason="Bad Request")]
CALLS.clear(); SLEEPS.clear()
try:
    extraction._request_with_retry("POST", KEYED, json={}, _sleep=_sleep)
    check("400 raises", False)
except extraction.GeminiHTTPError as exc:
    check("400: fails once, names the detail, key-free",
          len(CALLS) == 1 and "INVALID_ARGUMENT" in str(exc) and "SECRET" not in str(exc))

# --- 4. connection error quoting the keyed URL is re-raised clean ---------------------
conn_exc = requests.ConnectionError(f"HTTPSConnectionPool(host='g', port=443): Max retries exceeded with url: /v1beta/models/x:generateContent?key=AQ.SECRETSECRETSECRET (Caused by X)")
SCRIPT[:] = [conn_exc] * (extraction.MAX_RETRIES + 1)
CALLS.clear(); SLEEPS.clear()
try:
    extraction._request_with_retry("POST", KEYED, json={}, _sleep=_sleep)
    check("connection error raises", False)
except requests.RequestException as exc:
    check("connection error: retried, then re-raised without the key",
          len(CALLS) == extraction.MAX_RETRIES + 1 and "SECRET" not in str(exc) and "<redacted>" in str(exc))

# --- 5. batch_ingest: gemini_error result, quota flag, key-free ---------------------------
import batch_ingest
from pathlib import Path
import tempfile

def raising_extract(pdf_bytes, filename, api_key):
    raise extraction.GeminiHTTPError("429 RESOURCE_EXHAUSTED from Gemini: Your project has exceeded its monthly spending cap.",
                                     FakeResp(429, cap_body), "RESOURCE_EXHAUSTED", "monthly spending cap")

def leaking_extract(pdf_bytes, filename, api_key):
    raise requests.HTTPError(f"429 Client Error: Too Many Requests for url: {KEYED}")

class _DB:      # minimal backend: nothing is seen, nothing is written
    @staticmethod
    def seen_sha1(conn, sha1): return False

tmp = Path(tempfile.mkdtemp()) / "x.pdf"
tmp.write_bytes(b"%PDF-1.4 fake")
batch_ingest.extract_from_pdf = raising_extract
res = batch_ingest.process_pdf(tmp, None, "k", db=_DB)
check("process_pdf: quota refusal → gemini_error result with quota=True",
      (res.error or "").startswith("gemini_error") and res.quota is True and "spending cap" in res.error)
batch_ingest.extract_from_pdf = leaking_extract
res = batch_ingest.process_pdf(tmp, None, "k", db=_DB)
check("process_pdf: a raw HTTPError with the keyed URL is stored key-free",
      (res.error or "").startswith("gemini_error") and "SECRET" not in res.error and "<redacted>" in res.error
      and res.quota is False)

# --- 5b. thinking option: applied to every call; 400 on it → retried without, then off ----
extraction._THINKING_UNSUPPORTED = False
extraction.THINKING = "low"
check("thinking_config: low → thinkingLevel", extraction.thinking_config() == {"thinkingLevel": "low"})
check("thinking_config: off → budget 0", extraction.thinking_config("off") == {"thinkingBudget": 0})
check("thinking_config: 2048 → budget", extraction.thinking_config("2048") == {"thinkingBudget": 2048})
check("thinking_config: dynamic → nothing sent", extraction.thinking_config("dynamic") is None)
check("thinking_config: garbage → nothing sent", extraction.thinking_config("lots") is None)
SENT: list = []
def fake_request2(method, url, timeout=None, **kw):
    import copy
    SENT.append(copy.deepcopy(kw.get("json")))   # the payload dict is mutated on retry
    nxt = SCRIPT.pop(0) if SCRIPT else FakeResp(200, {})
    if isinstance(nxt, Exception):
        raise nxt
    return nxt
extraction.requests.request = fake_request2
SCRIPT[:] = [FakeResp(200, {"ok": 1})]; SENT.clear()
extraction.gemini_request(KEYED, {"contents": [], "generationConfig": {"temperature": 0}}, _sleep=_sleep)
check("gemini_request adds thinkingConfig to the payload",
      SENT and SENT[0]["generationConfig"].get("thinkingConfig") == {"thinkingLevel": "low"})
# the model rejects the option
SCRIPT[:] = [FakeResp(400, {"error": {"code": 400, "status": "INVALID_ARGUMENT",
                                      "message": "Invalid value at 'generation_config.thinking_config.thinking_level'"}},
                      reason="Bad Request"),
             FakeResp(200, {"ok": 2})]
SENT.clear()
r = extraction.gemini_request(KEYED, {"contents": [], "generationConfig": {"temperature": 0}}, _sleep=_sleep)
check("400 naming thinking → retried once without thinkingConfig, 200 returned",
      r is not None and r.status_code == 200 and len(SENT) == 2
      and "thinkingConfig" in SENT[0]["generationConfig"] and "thinkingConfig" not in SENT[1]["generationConfig"])
check("…and the option is off for the rest of the process", extraction._THINKING_UNSUPPORTED is True
      and extraction.thinking_config() is None)
SCRIPT[:] = [FakeResp(200, {"ok": 3})]; SENT.clear()
extraction.gemini_request(KEYED, {"contents": [], "generationConfig": {}}, _sleep=_sleep)
check("later calls send no thinkingConfig", "thinkingConfig" not in SENT[0]["generationConfig"])
extraction._THINKING_UNSUPPORTED = False
# an unrelated 400 is not retried
SCRIPT[:] = [FakeResp(400, {"error": {"code": 400, "status": "INVALID_ARGUMENT", "message": "bad schema"}}, reason="Bad Request")]
SENT.clear()
try:
    extraction.gemini_request(KEYED, {"contents": [], "generationConfig": {}}, _sleep=_sleep)
    check("unrelated 400 raises", False)
except extraction.GeminiHTTPError:
    check("unrelated 400: raised once, thinking option kept", len(SENT) == 1 and extraction._THINKING_UNSUPPORTED is False)
extraction.requests.request = fake_request

# --- 6. agdb.redact ------------------------------------------------------------------------
check("redact: ?key=", agdb.redact("x for url: https://g/v1?key=AQ.abc-def") == "x for url: https://g/v1?key=<redacted>")
check("redact: &api_key= keeps the rest of the query", agdb.redact("/works?search=q&api_key=oa123&per-page=5") == "/works?search=q&api_key=<redacted>&per-page=5")
check("redact: bearer token", agdb.redact("Authorization: Bearer abcdefghijklmnop") == "Authorization: Bearer <redacted>")
check("redact: case-insensitive, JSON-embedded", agdb.redact('{"e":"…?Key=SECRET"}') == '{"e":"…?Key=<redacted>"}')
check("redact: already redacted is left alone", agdb.redact("u?key=<redacted>&a=1") == "u?key=<redacted>&a=1")
check("redact: plain text untouched", agdb.redact("the key to the frontier") == "the key to the frontier")
check("redact: non-strings pass through", agdb.redact(None) is None and agdb.redact(5) == 5)

print()
print("all tests passed" if not fails else f"{len(fails)} FAILED: {fails}")
sys.exit(1 if fails else 0)