"""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")