test_agent_state.py

"""Unit tests for agdb crawler-state persistence and the one-shot URL reset.

No database needed: agdb.get_state/set_state are replaced with a dict-backed
fake, so these run anywhere (unlike test_e2e.py, which wants a local Postgres).

    python test_agent_state.py
"""
import sys
import tempfile
from pathlib import Path

import pdf_crawler
from agent import agdb

# --- dict-backed agent_state ------------------------------------------------

KV: dict = {}


def fake_get_state(conn, key, default=None):
    return KV.get(key, default)


def fake_set_state(conn, key, value):
    KV[key] = value


agdb.get_state = fake_get_state
agdb.set_state = fake_set_state

fails = []


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


def fresh_state() -> pdf_crawler.CrawlerState:
    return pdf_crawler.CrawlerState(Path(tempfile.mkdtemp()) / "state.json")


# --- save/load round-trip, including transient-failure counts ----------------

KV.clear()
st = fresh_state()
st.seen_urls.add("https://example.org/a.pdf")
st.seen_hashes.add("deadbeef")
st.failed["https://example.org/blocked.pdf"] = 2
agdb.save_crawler_state(None, st)

st2 = fresh_state()
st2.failed["https://example.org/blocked.pdf"] = 1   # older local count
agdb.load_crawler_state(None, st2)
check("seen_urls round-trip", st2.seen_urls == {"https://example.org/a.pdf"})
check("seen_hashes round-trip", st2.seen_hashes == {"deadbeef"})
check("failed counts persisted, max wins",
      st2.failed == {"https://example.org/blocked.pdf": 2})

# --- one-shot URL reset -------------------------------------------------------

KV.clear()
KV["crawler_state"] = {
    "seen_urls": ["https://example.org/a.pdf", "https://example.org/b.pdf"],
    "seen_hashes": ["deadbeef"],
    "failed": {"https://example.org/blocked.pdf": 2},
}
n = agdb.reset_crawler_urls_once(None)
check("reset reports cleared count", n == 2)
check("reset drops seen_urls", KV["crawler_state"]["seen_urls"] == [])
check("reset drops failed counts", "failed" not in KV["crawler_state"])
check("reset keeps content hashes",
      KV["crawler_state"]["seen_hashes"] == ["deadbeef"])
check("reset leaves a marker", bool(KV.get("crawler_url_reset_v1")))
check("reset runs only once", agdb.reset_crawler_urls_once(None) is None)

# --- reset on a fresh install (no crawler_state yet) --------------------------

KV.clear()
n = agdb.reset_crawler_urls_once(None)
check("fresh install: reset is a no-op with 0 cleared", n == 0)
check("fresh install: marker still set so it never fires again",
      bool(KV.get("crawler_url_reset_v1")))

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