Spaces:
Running
Running
File size: 3,734 Bytes
be40a40 | 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 92 93 94 95 96 97 98 99 100 101 102 103 | """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
|