Spaces:
Running
Running
| """The discovery queue: AtomicQueue mutations, atomic claim under contention, | |
| and the HTTP endpoints.""" | |
| import threading | |
| def _wiki_env(make_env, **kw): | |
| return make_env(WIKI_ENABLED="true", WIKI_DATASET="test-org/kb", **kw) | |
| def _ids_by_status(env): | |
| out = {} | |
| for r in env.read_model.records("queue"): | |
| out.setdefault(r.frontmatter.get("status"), []).append(r.frontmatter.get("id")) | |
| return out | |
| def test_add_dedup_bumps_ref_count(make_env): | |
| env = _wiki_env(make_env) | |
| q = env.queue | |
| assert q.add([{"id": "arxiv:1", "title": "One"}, {"id": "arxiv:2"}]) == {"added": 2, "updated": 0} | |
| assert q.add([{"id": "arxiv:1", "discovered_from": ["arxiv:9"]}]) == {"added": 0, "updated": 1} | |
| fm = {r.frontmatter["id"]: r.frontmatter for r in env.read_model.records("queue")} | |
| assert fm["arxiv:1"]["ref_count"] == 1 | |
| assert fm["arxiv:1"]["discovered_from"] == ["arxiv:9"] | |
| def test_claim_prefers_highest_ref_count_then_exhausts(make_env): | |
| env = _wiki_env(make_env) | |
| q = env.queue | |
| q.add([{"id": "arxiv:1"}, {"id": "arxiv:2"}]) | |
| q.add([{"id": "arxiv:2"}]) # arxiv:2 now ref_count 1 -> higher priority | |
| assert q.claim("reader-a").id == "arxiv:2" | |
| assert q.claim("reader-b").id == "arxiv:1" | |
| assert q.claim("reader-c") is None # frontier exhausted | |
| def test_skip_and_mark_processed(make_env): | |
| env = _wiki_env(make_env) | |
| q = env.queue | |
| q.add([{"id": "arxiv:1"}, {"id": "arxiv:2"}]) | |
| assert q.skip("reader-a", "arxiv:1", "out of scope") is True | |
| q.mark_processed("arxiv:2", "sources/arxiv-2.md") | |
| by_status = _ids_by_status(env) | |
| assert by_status["skipped"] == ["arxiv:1"] | |
| assert by_status["processed"] == ["arxiv:2"] | |
| def test_expired_lease_is_reaped(make_env): | |
| env = _wiki_env(make_env, QUEUE_CLAIM_LEASE_S="0") # lease expires immediately | |
| q = env.queue | |
| q.add([{"id": "arxiv:1"}]) | |
| claimed = q.claim("reader-a") | |
| assert claimed.id == "arxiv:1" | |
| # lease already expired -> reaped back to the frontier, claimable again | |
| assert q.reap() == 1 | |
| assert q.claim("reader-b").id == "arxiv:1" | |
| def test_claim_is_atomic_under_contention(make_env): | |
| env = _wiki_env(make_env) | |
| q = env.queue | |
| n = 8 | |
| q.add([{"id": f"arxiv:{i}"} for i in range(n)]) | |
| got, lock, barrier = [], threading.Lock(), threading.Barrier(n) | |
| def worker(name): | |
| barrier.wait() | |
| it = q.claim(name) | |
| if it: | |
| with lock: | |
| got.append(it.id) | |
| threads = [threading.Thread(target=worker, args=(f"a{i}",)) for i in range(n)] | |
| for t in threads: | |
| t.start() | |
| for t in threads: | |
| t.join() | |
| # every source claimed exactly once — no double-claim | |
| assert len(got) == n and len(set(got)) == n | |
| def test_queue_endpoints(make_env): | |
| env = _wiki_env(make_env) | |
| client = env.client | |
| assert client.post("/v1/queue:add", json={"items": [{"id": "arxiv:7", "title": "Seven"}]}).json() == { | |
| "added": 1, "updated": 0 | |
| } | |
| listing = client.get("/v1/queue").json() | |
| assert listing["count"] == 1 and listing["items"][0]["id"] == "arxiv:7" | |
| claim = client.post("/v1/queue:claim", json={"agent_id": "reader-a"}).json() | |
| assert claim["claimed"] is True and claim["id"] == "arxiv:7" and claim["lease_seconds"] > 0 | |
| skip = client.post("/v1/queue:skip", json={"agent_id": "reader-a", "id": "arxiv:7", "reason": "dup"}) | |
| assert skip.status_code == 200 | |
| assert client.get("/v1/queue").json()["by_status"].get("skipped") == 1 | |
| def test_queue_404_when_wiki_disabled(env): | |
| # the default env fixture has wiki mode off | |
| assert env.client.post("/v1/queue:add", json={"items": []}).status_code == 404 | |
| assert env.client.get("/v1/queue").status_code == 404 | |