Spaces:
Running
Running
File size: 2,727 Bytes
2f3aa0d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """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")
|