feat: controller-owned allocator API on control-plane DB (Closes #600)
Add gitea_allocate_next_work and allocator_service routing policy on top of the #613 ControlPlaneDB substrate. Workers get atomic assign+lease results (or wait/no_safe_work/terminal-path outcomes) without self-selecting work via file locks or comment-only leases. #612 remains downstream. Closes #600
This commit is contained in:
@@ -0,0 +1,366 @@
|
||||
"""Tests for controller-owned allocator (#600) on control-plane DB (#613)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from allocator_service import (
|
||||
OUTCOME_ASSIGNED,
|
||||
OUTCOME_BLOCKED_TERMINAL,
|
||||
OUTCOME_NO_SAFE,
|
||||
OUTCOME_PREVIEW,
|
||||
OUTCOME_WAIT,
|
||||
WorkCandidate,
|
||||
allocate_next_work,
|
||||
candidate_from_dict,
|
||||
classify_skip,
|
||||
expected_role_for_candidate,
|
||||
)
|
||||
from control_plane_db import ControlPlaneDB, InvalidWorkKindError
|
||||
|
||||
|
||||
class AllocatorServiceTest(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self._tmp.name, "cp.sqlite3")
|
||||
self.db = ControlPlaneDB(self.db_path)
|
||||
|
||||
def tearDown(self) -> None:
|
||||
self._tmp.cleanup()
|
||||
|
||||
def _alloc(self, **kwargs):
|
||||
defaults = dict(
|
||||
db=self.db,
|
||||
session_id="s-test",
|
||||
role="author",
|
||||
remote="prgs",
|
||||
org="org",
|
||||
repo="repo",
|
||||
candidates=[],
|
||||
apply=False,
|
||||
profile_name="prgs-author",
|
||||
username="jcwalker3",
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return allocate_next_work(**defaults)
|
||||
|
||||
def test_selects_ready_issue_for_author(self) -> None:
|
||||
cands = [
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=612,
|
||||
labels=("status:ready",),
|
||||
title="bridge",
|
||||
dependency_unmet=True,
|
||||
dependency_reason="downstream of #600",
|
||||
),
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=600,
|
||||
labels=("status:ready", "type:feature"),
|
||||
title="allocator",
|
||||
priority=50,
|
||||
),
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=601,
|
||||
labels=("status:blocked",),
|
||||
title="blocked",
|
||||
blocked=True,
|
||||
),
|
||||
]
|
||||
res = self._alloc(candidates=cands, apply=False)
|
||||
self.assertTrue(res["success"])
|
||||
self.assertEqual(res["outcome"], OUTCOME_PREVIEW)
|
||||
self.assertEqual(res["selected"]["number"], 600)
|
||||
self.assertEqual(res["selected"]["kind"], "issue")
|
||||
# When 600 is highest priority valid work, lower-priority blocked/dep
|
||||
# candidates are not visited. Prove they are skipped when they sort first.
|
||||
res2 = self._alloc(
|
||||
candidates=[
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=612,
|
||||
labels=("status:ready",),
|
||||
priority=99,
|
||||
dependency_unmet=True,
|
||||
dependency_reason="downstream of #600",
|
||||
),
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=601,
|
||||
labels=("status:blocked",),
|
||||
priority=98,
|
||||
blocked=True,
|
||||
),
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=600,
|
||||
labels=("status:ready",),
|
||||
priority=1,
|
||||
),
|
||||
],
|
||||
apply=False,
|
||||
)
|
||||
skipped_nums = {s["number"] for s in res2["skipped"]}
|
||||
self.assertIn(612, skipped_nums)
|
||||
self.assertIn(601, skipped_nums)
|
||||
self.assertEqual(res2["selected"]["number"], 600)
|
||||
self.assertIsNone(res["assignment"])
|
||||
self.assertFalse(res["file_lock_only"])
|
||||
self.assertFalse(res["comment_lease_only"])
|
||||
|
||||
def test_atomic_assign_and_lease_on_apply(self) -> None:
|
||||
cands = [
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=600,
|
||||
labels=("status:ready",),
|
||||
priority=10,
|
||||
)
|
||||
]
|
||||
res = self._alloc(candidates=cands, apply=True, session_id="s-a")
|
||||
self.assertEqual(res["outcome"], OUTCOME_ASSIGNED)
|
||||
asn = res["assignment"]
|
||||
self.assertEqual(asn["outcome"], "assigned")
|
||||
self.assertIsNotNone(asn["assignment_id"])
|
||||
self.assertIsNotNone(asn["lease_id"])
|
||||
self.assertEqual(asn["work_number"], 600)
|
||||
self.assertIn("implement", asn["allowed_actions"])
|
||||
self.assertIn("merge", asn["forbidden_actions"])
|
||||
proof = res["lease_proof"]
|
||||
self.assertEqual(proof["source"], "control_plane_db.assign_and_lease")
|
||||
|
||||
def test_concurrent_allocators_no_double_assign(self) -> None:
|
||||
cand = WorkCandidate(
|
||||
kind="pr",
|
||||
number=100,
|
||||
head_sha="a" * 40,
|
||||
priority=10,
|
||||
)
|
||||
# Seed work item so both race the same target
|
||||
self.db.upsert_work_item(
|
||||
remote="prgs",
|
||||
org="org",
|
||||
repo="repo",
|
||||
kind="pr",
|
||||
number=100,
|
||||
current_head_sha="a" * 40,
|
||||
)
|
||||
results = []
|
||||
lock = threading.Lock()
|
||||
|
||||
def worker(sid: str):
|
||||
r = allocate_next_work(
|
||||
self.db,
|
||||
session_id=sid,
|
||||
role="reviewer",
|
||||
remote="prgs",
|
||||
org="org",
|
||||
repo="repo",
|
||||
candidates=[cand],
|
||||
apply=True,
|
||||
profile_name="prgs-reviewer",
|
||||
)
|
||||
with lock:
|
||||
results.append(r)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
futs = [pool.submit(worker, f"s-{i}") for i in range(2)]
|
||||
for f in as_completed(futs):
|
||||
f.result()
|
||||
|
||||
outcomes = [r["outcome"] for r in results]
|
||||
self.assertEqual(outcomes.count(OUTCOME_ASSIGNED), 1, outcomes)
|
||||
self.assertEqual(outcomes.count(OUTCOME_WAIT), 1, outcomes)
|
||||
|
||||
def test_blocked_and_dependency_skipped(self) -> None:
|
||||
cands = [
|
||||
WorkCandidate(kind="issue", number=1, blocked=True, labels=("status:blocked",)),
|
||||
WorkCandidate(
|
||||
kind="issue",
|
||||
number=612,
|
||||
labels=("status:ready",),
|
||||
dependency_unmet=True,
|
||||
dependency_reason="downstream of #600",
|
||||
),
|
||||
]
|
||||
res = self._alloc(candidates=cands, apply=True, role="author")
|
||||
self.assertEqual(res["outcome"], OUTCOME_NO_SAFE)
|
||||
self.assertIsNone(res["selected"])
|
||||
reasons = " ".join(s["reason"] for s in res["skipped"])
|
||||
self.assertIn("blocked", reasons.lower())
|
||||
self.assertIn("600", reasons)
|
||||
|
||||
def test_already_leased_returns_wait(self) -> None:
|
||||
self.db.upsert_session(session_id="owner", role="author")
|
||||
self.db.assign_and_lease(
|
||||
session_id="owner",
|
||||
role="author",
|
||||
remote="prgs",
|
||||
org="org",
|
||||
repo="repo",
|
||||
kind="issue",
|
||||
number=50,
|
||||
)
|
||||
res = self._alloc(
|
||||
session_id="other",
|
||||
candidates=[
|
||||
WorkCandidate(kind="issue", number=50, labels=("status:ready",))
|
||||
],
|
||||
apply=True,
|
||||
)
|
||||
self.assertEqual(res["outcome"], OUTCOME_WAIT)
|
||||
self.assertEqual(res["owner_session_id"], "owner")
|
||||
|
||||
def test_role_ineligible_skipped(self) -> None:
|
||||
# PR with current-head REQUEST_CHANGES expects author, not reviewer.
|
||||
cand = WorkCandidate(
|
||||
kind="pr",
|
||||
number=9,
|
||||
head_sha="b" * 40,
|
||||
request_changes_current_head=True,
|
||||
)
|
||||
self.assertEqual(expected_role_for_candidate(cand), "author")
|
||||
res = self._alloc(
|
||||
role="reviewer",
|
||||
profile_name="prgs-reviewer",
|
||||
candidates=[cand],
|
||||
apply=False,
|
||||
)
|
||||
self.assertEqual(res["outcome"], OUTCOME_NO_SAFE)
|
||||
self.assertTrue(any("expects role" in s["reason"] for s in res["skipped"]))
|
||||
|
||||
def test_terminal_lock_blocks_downstream_reviewer_prs(self) -> None:
|
||||
self.db.set_terminal_lock(
|
||||
remote="prgs",
|
||||
org="org",
|
||||
repo="repo",
|
||||
terminal_pr=10,
|
||||
decision="request_changes",
|
||||
status="active",
|
||||
)
|
||||
cands = [
|
||||
WorkCandidate(kind="pr", number=11, head_sha="c" * 40, priority=5),
|
||||
WorkCandidate(kind="pr", number=10, head_sha="d" * 40, priority=1),
|
||||
]
|
||||
res = self._alloc(
|
||||
role="reviewer",
|
||||
profile_name="prgs-reviewer",
|
||||
candidates=cands,
|
||||
apply=False,
|
||||
)
|
||||
# Terminal PR #10 is selectable; #11 skipped for terminal path.
|
||||
self.assertEqual(res["selected"]["number"], 10)
|
||||
self.assertTrue(
|
||||
any(
|
||||
s["number"] == 11 and "terminal-review lock" in s["reason"]
|
||||
for s in res["skipped"]
|
||||
)
|
||||
)
|
||||
|
||||
def test_terminal_lock_blocks_all_when_only_downstream(self) -> None:
|
||||
self.db.set_terminal_lock(
|
||||
remote="prgs",
|
||||
org="org",
|
||||
repo="repo",
|
||||
terminal_pr=10,
|
||||
decision="request_changes",
|
||||
status="active",
|
||||
)
|
||||
res = self._alloc(
|
||||
role="reviewer",
|
||||
profile_name="prgs-reviewer",
|
||||
candidates=[
|
||||
WorkCandidate(kind="pr", number=99, head_sha="e" * 40),
|
||||
],
|
||||
apply=False,
|
||||
)
|
||||
self.assertEqual(res["outcome"], OUTCOME_BLOCKED_TERMINAL)
|
||||
|
||||
def test_no_work_structured(self) -> None:
|
||||
res = self._alloc(candidates=[], apply=True)
|
||||
self.assertEqual(res["outcome"], OUTCOME_NO_SAFE)
|
||||
self.assertIsNone(res["selected"])
|
||||
self.assertTrue(res["reasons"])
|
||||
|
||||
def test_rejects_incident_kind(self) -> None:
|
||||
with self.assertRaises(InvalidWorkKindError):
|
||||
WorkCandidate(kind="sentry_incident", number=1)
|
||||
|
||||
def test_612_downstream_marker_in_result(self) -> None:
|
||||
res = self._alloc(
|
||||
candidates=[
|
||||
WorkCandidate(kind="issue", number=600, labels=("status:ready",))
|
||||
],
|
||||
apply=True,
|
||||
)
|
||||
self.assertIn("612", res.get("downstream_note", ""))
|
||||
|
||||
def test_substrate_not_file_or_comment_lease(self) -> None:
|
||||
res = self._alloc(
|
||||
candidates=[
|
||||
WorkCandidate(kind="issue", number=1, labels=("status:ready",))
|
||||
],
|
||||
apply=True,
|
||||
)
|
||||
self.assertEqual(res["substrate"], "control_plane_db")
|
||||
self.assertFalse(res["file_lock_only"])
|
||||
self.assertFalse(res["comment_lease_only"])
|
||||
self.assertEqual(
|
||||
res["lease_proof"]["source"], "control_plane_db.assign_and_lease"
|
||||
)
|
||||
|
||||
def test_candidate_from_dict(self) -> None:
|
||||
c = candidate_from_dict(
|
||||
{
|
||||
"kind": "pr",
|
||||
"number": 7,
|
||||
"head_sha": "f" * 40,
|
||||
"approval_on_current_head": True,
|
||||
"mergeable": True,
|
||||
}
|
||||
)
|
||||
self.assertEqual(expected_role_for_candidate(c), "merger")
|
||||
|
||||
def test_merger_gets_clean_approval(self) -> None:
|
||||
c = WorkCandidate(
|
||||
kind="pr",
|
||||
number=3,
|
||||
head_sha="1" * 40,
|
||||
approval_on_current_head=True,
|
||||
mergeable=True,
|
||||
priority=100,
|
||||
)
|
||||
res = self._alloc(
|
||||
role="merger",
|
||||
profile_name="prgs-merger",
|
||||
candidates=[c],
|
||||
apply=True,
|
||||
session_id="merger-1",
|
||||
)
|
||||
self.assertEqual(res["outcome"], OUTCOME_ASSIGNED)
|
||||
self.assertEqual(res["assignment"]["work_number"], 3)
|
||||
self.assertIn("merge", res["assignment"]["allowed_actions"])
|
||||
|
||||
def test_db_unavailable_fails_closed(self) -> None:
|
||||
res = allocate_next_work(
|
||||
None, # type: ignore[arg-type]
|
||||
session_id="x",
|
||||
role="author",
|
||||
remote="prgs",
|
||||
org="o",
|
||||
repo="r",
|
||||
candidates=[],
|
||||
apply=True,
|
||||
)
|
||||
self.assertFalse(res["success"])
|
||||
self.assertIn("unavailable", res["reasons"][0].lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user