"""Tests for control-plane DB substrate (#613).""" from __future__ import annotations import os import tempfile import threading import unittest from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import timedelta from control_plane_db import ( ControlPlaneDB, InvalidWorkKindError, LeaseRequiredError, WORK_KINDS, _ts, _utc_now, ) class ControlPlaneDBTest(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 test_schema_and_architecture_meta(self) -> None: import sqlite3 conn = sqlite3.connect(self.db_path) try: rows = dict(conn.execute("SELECT key, value FROM schema_meta").fetchall()) finally: conn.close() self.assertEqual(rows["schema_version"], "2") self.assertIn("DB coordinates", rows["architecture"]) self.assertIn("bridge", rows["architecture"].lower()) def test_rejects_raw_incident_as_work_kind(self) -> None: with self.assertRaises(InvalidWorkKindError): self.db.upsert_work_item( remote="prgs", org="org", repo="repo", kind="sentry_incident", number=1, ) with self.assertRaises(InvalidWorkKindError): self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="org", repo="repo", kind="glitchtip_incident", number=9, ) self.assertEqual(WORK_KINDS, frozenset({"issue", "pr"})) def test_atomic_assign_and_lease_fields(self) -> None: self.db.upsert_session(session_id="s-a", role="author", profile="prgs-author") result = self.db.assign_and_lease( session_id="s-a", role="author", remote="prgs", org="Scaled-Tech-Consulting", repo="Gitea-Tools", kind="issue", number=613, expected_head_sha="abc123", allowed_actions=("implement", "comment"), forbidden_actions=("approve", "merge"), ) self.assertEqual(result.outcome, "assigned") self.assertIsNotNone(result.assignment_id) self.assertIsNotNone(result.lease_id) self.assertEqual(result.role, "author") self.assertEqual(result.work_kind, "issue") self.assertEqual(result.work_number, 613) self.assertEqual(result.expected_head_sha, "abc123") self.assertIn("implement", result.allowed_actions) self.assertIn("merge", result.forbidden_actions) self.assertIsNotNone(result.expires_at) def test_second_session_waits_on_foreign_lease(self) -> None: self.db.upsert_session(session_id="s1", role="author") self.db.upsert_session(session_id="s2", role="author") first = self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="pr", number=100, expected_head_sha="deadbeef", ) self.assertEqual(first.outcome, "assigned") second = self.db.assign_and_lease( session_id="s2", role="author", remote="prgs", org="o", repo="r", kind="pr", number=100, expected_head_sha="deadbeef", ) self.assertEqual(second.outcome, "wait") self.assertEqual(second.owner_session_id, "s1") def test_owner_resume_refreshes_lease(self) -> None: self.db.upsert_session(session_id="s1", role="reviewer") a = self.db.assign_and_lease( session_id="s1", role="reviewer", remote="prgs", org="o", repo="r", kind="pr", number=50, expected_head_sha="head-50", ) b = self.db.assign_and_lease( session_id="s1", role="reviewer", remote="prgs", org="o", repo="r", kind="pr", number=50, expected_head_sha="head-50", ) self.assertEqual(b.outcome, "assigned") self.assertEqual(b.lease_id, a.lease_id) self.assertIn("owner-resume", b.reason) def test_require_valid_assignment_gates_mutations(self) -> None: self.db.upsert_session(session_id="s1", role="author") self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="issue", number=7, allowed_actions=("implement",), forbidden_actions=("merge",), ) proof = self.db.require_valid_assignment( session_id="s1", remote="prgs", org="o", repo="r", kind="issue", number=7, action="implement", ) self.assertEqual(proof["session_id"], "s1") with self.assertRaises(LeaseRequiredError): self.db.require_valid_assignment( session_id="s1", remote="prgs", org="o", repo="r", kind="issue", number=7, action="merge", ) with self.assertRaises(LeaseRequiredError): self.db.require_valid_assignment( session_id="s-other", remote="prgs", org="o", repo="r", kind="issue", number=7, action="implement", ) def test_expired_lease_allows_reassign(self) -> None: self.db.upsert_session(session_id="s1", role="author") self.db.upsert_session(session_id="s2", role="author") past = _utc_now() - timedelta(hours=1) # Create lease already expired by using negative TTL edge via direct assign then expire assigned = self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="issue", number=3, lease_ttl_seconds=1, ) self.assertEqual(assigned.outcome, "assigned") # Force expiry in DB import sqlite3 conn = sqlite3.connect(self.db_path) try: conn.execute( "UPDATE leases SET expires_at = ? WHERE lease_id = ?", (_ts(past), assigned.lease_id), ) conn.commit() finally: conn.close() n = self.db.expire_stale_leases() self.assertGreaterEqual(n, 1) second = self.db.assign_and_lease( session_id="s2", role="author", remote="prgs", org="o", repo="r", kind="issue", number=3, ) self.assertEqual(second.outcome, "assigned") self.assertEqual(second.session_id, "s2") def test_merged_work_never_assigned(self) -> None: self.db.upsert_session(session_id="s1", role="merger") self.db.upsert_work_item( remote="prgs", org="o", repo="r", kind="pr", number=99, state="merged", ) result = self.db.assign_and_lease( session_id="s1", role="merger", remote="prgs", org="o", repo="r", kind="pr", number=99, expected_head_sha="merged-head", ) self.assertEqual(result.outcome, "no_safe_work") def test_four_concurrent_sessions_unique_assignments(self) -> None: """Four concurrent assigners on four different issues — all succeed uniquely. Also two concurrent assigners on the *same* issue: at most one assigned. """ for i in range(4): self.db.upsert_session(session_id=f"sess-{i}", role="author") def claim_unique(i: int): return self.db.assign_and_lease( session_id=f"sess-{i}", role="author", remote="prgs", org="o", repo="r", kind="issue", number=1000 + i, ) with ThreadPoolExecutor(max_workers=4) as pool: results = [f.result() for f in as_completed([pool.submit(claim_unique, i) for i in range(4)])] self.assertEqual({r.outcome for r in results}, {"assigned"}) numbers = sorted(r.work_number for r in results) self.assertEqual(numbers, [1000, 1001, 1002, 1003]) # Contention on one item self.db.upsert_session(session_id="c1", role="author") self.db.upsert_session(session_id="c2", role="author") self.db.upsert_session(session_id="c3", role="author") self.db.upsert_session(session_id="c4", role="author") barrier = threading.Barrier(4) outcomes: list[str] = [] lock = threading.Lock() def contend(sid: str) -> None: barrier.wait() res = self.db.assign_and_lease( session_id=sid, role="author", remote="prgs", org="o", repo="r", kind="issue", number=7777, ) with lock: outcomes.append(res.outcome) threads = [threading.Thread(target=contend, args=(f"c{i}",)) for i in range(1, 5)] for t in threads: t.start() for t in threads: t.join() self.assertEqual(outcomes.count("assigned"), 1) self.assertEqual(outcomes.count("wait"), 3) def test_terminal_lock_index(self) -> None: self.db.set_terminal_lock( remote="prgs", org="o", repo="r", terminal_pr=332, review_id="rev-1", decision="approve", ) row = self.db.get_active_terminal_lock(remote="prgs", org="o", repo="r") self.assertIsNotNone(row) assert row is not None self.assertEqual(row["terminal_pr"], 332) self.assertEqual(row["status"], "active") def test_incident_links_not_work_items(self) -> None: link = self.db.upsert_incident_link( provider="sentry", provider_base_url="https://sentry.prgs.cc", provider_org="prgs", provider_project="gitea-tools-mcp", provider_issue_id="12345", gitea_org="Scaled-Tech-Consulting", gitea_repo="Gitea-Tools", gitea_issue_number=9001, fingerprint="fp-1", ) self.assertEqual(link["gitea_issue_number"], 9001) found = self.db.get_incident_link_for_gitea_issue( gitea_org="Scaled-Tech-Consulting", gitea_repo="Gitea-Tools", gitea_issue_number=9001, ) self.assertIsNotNone(found) # Linking does not create a work_item of incident kind with self.assertRaises(InvalidWorkKindError): self.db.upsert_work_item( remote="prgs", org="Scaled-Tech-Consulting", repo="Gitea-Tools", kind="sentry", number=12345, ) def test_heartbeat_and_release(self) -> None: self.db.upsert_session(session_id="s1", role="author") a = self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="issue", number=1, ) hb = self.db.heartbeat_lease(a.lease_id, session_id="s1") self.assertEqual(hb["lease_id"], a.lease_id) self.db.release_lease(a.lease_id, session_id="s1") # After release another session can claim self.db.upsert_session(session_id="s2", role="author") b = self.db.assign_and_lease( session_id="s2", role="author", remote="prgs", org="o", repo="r", kind="issue", number=1, ) self.assertEqual(b.outcome, "assigned") self.assertEqual(b.session_id, "s2") def test_require_valid_assignment_rejects_stale_head(self) -> None: """Assignment must not authorize mutations after work-item head drifts.""" self.db.upsert_session(session_id="s1", role="author") self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="pr", number=42, expected_head_sha="head-v1", allowed_actions=("implement",), ) # Head drifts after assignment self.db.upsert_work_item( remote="prgs", org="o", repo="r", kind="pr", number=42, current_head_sha="head-v2", ) with self.assertRaises(LeaseRequiredError) as ctx: self.db.require_valid_assignment( session_id="s1", remote="prgs", org="o", repo="r", kind="pr", number=42, action="implement", ) self.assertIn("stale head", str(ctx.exception).lower()) def test_require_valid_assignment_rejects_terminal_state(self) -> None: """Assignment must not authorize mutations after work item is merged/closed.""" self.db.upsert_session(session_id="s1", role="author") self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="pr", number=55, expected_head_sha="abc", allowed_actions=("implement",), ) self.db.upsert_work_item( remote="prgs", org="o", repo="r", kind="pr", number=55, state="merged", current_head_sha="abc", ) with self.assertRaises(LeaseRequiredError) as ctx: self.db.require_valid_assignment( session_id="s1", remote="prgs", org="o", repo="r", kind="pr", number=55, action="implement", ) self.assertIn("terminal", str(ctx.exception).lower()) self.db.upsert_work_item( remote="prgs", org="o", repo="r", kind="issue", number=56, state="open", ) self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="issue", number=56, allowed_actions=("implement",), ) self.db.upsert_work_item( remote="prgs", org="o", repo="r", kind="issue", number=56, state="closed", ) with self.assertRaises(LeaseRequiredError): self.db.require_valid_assignment( session_id="s1", remote="prgs", org="o", repo="r", kind="issue", number=56, action="implement", ) def test_pr_assignment_requires_expected_head_sha(self) -> None: """PR assign and mutation must fail closed without a head pin.""" self.db.upsert_session(session_id="s1", role="author") with self.assertRaises(LeaseRequiredError) as ctx: self.db.assign_and_lease( session_id="s1", role="author", remote="prgs", org="o", repo="r", kind="pr", number=88, expected_head_sha=None, allowed_actions=("implement",), ) self.assertIn("expected_head_sha", str(ctx.exception)) # Legacy path: force an unpinned PR assignment into the DB, then # populate head and prove mutation is still rejected. import sqlite3 self.db.upsert_work_item( remote="prgs", org="o", repo="r", kind="pr", number=89, current_head_sha=None, ) conn = sqlite3.connect(self.db_path) try: wid = conn.execute( "SELECT work_item_id FROM work_items WHERE kind='pr' AND number=89" ).fetchone()[0] conn.execute( """ INSERT INTO sessions(session_id, role, started_at, last_heartbeat_at, status) VALUES ('legacy', 'author', '2020-01-01T00:00:00Z', '2020-01-01T00:00:00Z', 'active') """ ) conn.execute( """ INSERT INTO leases( lease_id, work_item_id, session_id, role, phase, expires_at, heartbeat_at, status ) VALUES ( 'lease-legacy', ?, 'legacy', 'author', 'claimed', '2099-01-01T00:00:00Z', '2020-01-01T00:00:00Z', 'active' ) """, (wid,), ) conn.execute( """ INSERT INTO assignments( assignment_id, work_item_id, session_id, lease_id, allowed_actions, forbidden_actions, expected_head_sha, role, status, created_at ) VALUES ( 'asn-legacy', ?, 'legacy', 'lease-legacy', '["implement"]', '["merge"]', NULL, 'author', 'active', '2020-01-01T00:00:00Z' ) """, (wid,), ) conn.commit() finally: conn.close() self.db.upsert_work_item( remote="prgs", org="o", repo="r", kind="pr", number=89, current_head_sha="populated-head", ) with self.assertRaises(LeaseRequiredError) as ctx2: self.db.require_valid_assignment( session_id="legacy", remote="prgs", org="o", repo="r", kind="pr", number=89, action="implement", ) self.assertIn("expected_head_sha pin", str(ctx2.exception)) def test_migrate_duplicate_null_scope_incident_links(self) -> None: """Legacy NULL-scope duplicates must migrate without UNIQUE crash.""" import sqlite3 from control_plane_db import ControlPlaneDB, ControlPlaneError # Build a v1-like table with nullable scope columns and insert dups # that collapse under normalization, then open ControlPlaneDB on it. path = os.path.join(self._tmp.name, "legacy_dups.sqlite3") conn = sqlite3.connect(path) try: conn.executescript( """ CREATE TABLE schema_meta (key TEXT PRIMARY KEY, value TEXT NOT NULL); CREATE TABLE incident_links ( link_id INTEGER PRIMARY KEY AUTOINCREMENT, provider TEXT NOT NULL, provider_base_url TEXT, provider_org TEXT, provider_project TEXT, provider_issue_id TEXT NOT NULL, provider_short_id TEXT, provider_permalink TEXT, fingerprint TEXT, gitea_org TEXT NOT NULL, gitea_repo TEXT NOT NULL, gitea_issue_number INTEGER NOT NULL, linked_pr_numbers TEXT, first_seen TEXT, last_seen TEXT, event_count INTEGER, status TEXT NOT NULL DEFAULT 'open', release_resolved_at TEXT, last_sync_at TEXT, UNIQUE ( provider, provider_base_url, provider_org, provider_project, provider_issue_id ) ); INSERT INTO incident_links( provider, provider_base_url, provider_org, provider_project, provider_issue_id, gitea_org, gitea_repo, gitea_issue_number, status ) VALUES ('sentry', NULL, NULL, NULL, 'dup-1', 'org', 'repo', 10, 'open'), ('sentry', NULL, NULL, NULL, 'dup-1', 'org', 'repo', 10, 'open'); """ ) # SQLite allows two NULL-scope rows with same provider/issue under UNIQUE. n = conn.execute("SELECT COUNT(*) FROM incident_links").fetchone()[0] self.assertEqual(n, 2) conn.commit() finally: conn.close() db = ControlPlaneDB(path) conn2 = sqlite3.connect(path) try: n2 = conn2.execute("SELECT COUNT(*) FROM incident_links").fetchone()[0] rows = conn2.execute( "SELECT provider_base_url, provider_org, provider_project, gitea_issue_number " "FROM incident_links" ).fetchall() finally: conn2.close() self.assertEqual(n2, 1) self.assertEqual(rows[0][0], "") self.assertEqual(rows[0][1], "") self.assertEqual(rows[0][2], "") self.assertEqual(rows[0][3], 10) # Touch to silence unused import in type checkers if needed self.assertTrue(issubclass(ControlPlaneError, Exception)) del db def test_migrate_conflicting_duplicate_incident_links_fails_closed(self) -> None: """Conflicting Gitea targets for the same provider key must fail closed.""" import sqlite3 from control_plane_db import ControlPlaneDB, ControlPlaneError path = os.path.join(self._tmp.name, "legacy_conflict.sqlite3") conn = sqlite3.connect(path) try: conn.executescript( """ CREATE TABLE incident_links ( link_id INTEGER PRIMARY KEY AUTOINCREMENT, provider TEXT NOT NULL, provider_base_url TEXT, provider_org TEXT, provider_project TEXT, provider_issue_id TEXT NOT NULL, gitea_org TEXT NOT NULL, gitea_repo TEXT NOT NULL, gitea_issue_number INTEGER NOT NULL, status TEXT NOT NULL DEFAULT 'open', UNIQUE ( provider, provider_base_url, provider_org, provider_project, provider_issue_id ) ); INSERT INTO incident_links( provider, provider_base_url, provider_org, provider_project, provider_issue_id, gitea_org, gitea_repo, gitea_issue_number ) VALUES ('sentry', NULL, NULL, NULL, 'dup-c', 'org', 'repo', 1), ('sentry', NULL, NULL, NULL, 'dup-c', 'org', 'repo', 2); """ ) conn.commit() finally: conn.close() with self.assertRaises(ControlPlaneError) as ctx: ControlPlaneDB(path) self.assertIn("conflicting", str(ctx.exception).lower()) def test_incident_links_minimal_upsert_is_canonical(self) -> None: """Repeated minimal upserts must update one row (NULL-safe uniqueness).""" a = self.db.upsert_incident_link( provider="sentry", provider_issue_id="inc-1", gitea_org="org", gitea_repo="repo", gitea_issue_number=1, ) b = self.db.upsert_incident_link( provider="sentry", provider_issue_id="inc-1", gitea_org="org", gitea_repo="repo", gitea_issue_number=2, # omit optional scope fields again ) c = self.db.upsert_incident_link( provider="sentry", provider_issue_id="inc-1", gitea_org="org", gitea_repo="repo", gitea_issue_number=3, provider_base_url=None, provider_org="", provider_project=" ", ) self.assertEqual(a["link_id"], b["link_id"]) self.assertEqual(b["link_id"], c["link_id"]) self.assertEqual(c["gitea_issue_number"], 3) self.assertEqual(c["provider_base_url"], "") self.assertEqual(c["provider_org"], "") self.assertEqual(c["provider_project"], "") import sqlite3 conn = sqlite3.connect(self.db_path) try: n = conn.execute( "SELECT COUNT(*) FROM incident_links WHERE provider_issue_id = ?", ("inc-1",), ).fetchone()[0] finally: conn.close() self.assertEqual(n, 1) if __name__ == "__main__": unittest.main()