from __future__ import annotations import sqlite3 import tempfile import unittest from datetime import datetime, timedelta, timezone from pathlib import Path from ai_triage.models import ( FailureEvent, MonitorTarget, WorkflowJob, WorkflowRun, WorkflowStep, ) from ai_triage.state import StateStore from ai_triage import state as state_module NOW = datetime(2026, 7, 18, 2, 0, tzinfo=timezone.utc) TARGET = MonitorTarget( id="dedicated", enabled=True, workflow="scenario.yaml", branch="main", ) OTHER_TARGET = MonitorTarget( id="nextgen", enabled=True, workflow="nextgen-scenario.yaml", branch="main", ) TARGET_IDS = (TARGET.id,) def failure_event( *, target: MonitorTarget = TARGET, run_id: int = 10, attempt: int = 2, job_id: int | None = 30, ) -> FailureEvent: run = WorkflowRun( id=run_id, status="completed", conclusion="failure", run_attempt=attempt, html_url=f"https://github.example/runs/{run_id}", created_at=NOW - timedelta(hours=1), updated_at=NOW, ) job = None if job_id is not None: job = WorkflowJob( id=job_id, name="gcp (case.yaml)", status="completed", conclusion="failure", steps=( WorkflowStep(13, "Run scenario cases", "completed", "failure"), ), html_url=f"https://github.example/jobs/{job_id}", ) return FailureEvent.create(target, run, job) class StateStoreTests(unittest.TestCase): def setUp(self) -> None: self.directory = tempfile.TemporaryDirectory() self.path = Path(self.directory.name, "state", "events.sqlite3") def tearDown(self) -> None: self.directory.cleanup() def test_target_baseline_is_created_once_and_persists(self) -> None: with StateStore(self.path) as store: self.assertEqual(NOW, store.initialize_target("dedicated", NOW)) self.assertEqual( NOW, store.initialize_target("dedicated", NOW + timedelta(hours=1)), ) with StateStore(self.path) as reopened: self.assertEqual( NOW, reopened.initialize_target("dedicated", NOW + timedelta(days=1)), ) def test_startup_migrates_existing_database_for_notification_progress(self) -> None: self.path.parent.mkdir(parents=True) legacy_schema = state_module._SCHEMA.replace( " notification_progress_json TEXT,\n", "", ) with sqlite3.connect(self.path) as connection: connection.executescript(legacy_schema) with StateStore(self.path) as store: columns = { str(row["name"]) for row in store.connection.execute( "PRAGMA table_info(events)" ).fetchall() } self.assertIn("notification_progress_json", columns) def test_discovery_is_unique_per_run_attempt_and_job(self) -> None: with StateStore(self.path) as store: event = failure_event(run_id=10, attempt=2, job_id=30) self.assertTrue(store.discover(event, NOW)) self.assertFalse(store.discover(event, NOW)) self.assertTrue( store.discover( failure_event(run_id=10, attempt=3, job_id=31), NOW ) ) def test_claim_lease_prevents_parallel_work_and_expiry_recovers(self) -> None: event = failure_event() with StateStore(self.path) as first, StateStore(self.path) as second: first.discover(event, NOW) claimed = first.claim_next( "worker-a", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) self.assertEqual("triage", claimed.stage) self.assertIsNone( second.claim_next( "worker-b", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) ) recovered = second.claim_next( "worker-b", NOW + timedelta(seconds=61), target_ids=TARGET_IDS, lease_seconds=60, ) self.assertEqual(event.event_id, recovered.event.event_id) self.assertEqual("triage", recovered.stage) def test_claim_next_only_returns_events_for_allowed_targets(self) -> None: other_event = failure_event( target=OTHER_TARGET, run_id=9, job_id=29, ) target_event = failure_event(run_id=10, job_id=30) with StateStore(self.path) as store: store.discover(other_event, NOW) store.discover(target_event, NOW) claimed = store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) other_record = store.get_event(other_event.event_id) assert claimed is not None assert other_record is not None self.assertEqual(target_event.event_id, claimed.event.event_id) self.assertEqual("discovered", other_record.status) def test_completed_stages_resume_without_repeating_triage(self) -> None: event = failure_event() report = {"event_id": event.event_id, "incidents": []} notification = {"event_id": event.event_id, "incidents": []} with StateStore(self.path) as store: store.discover(event, NOW) triage_claim = store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) store.complete_triage( event.event_id, "worker", report, notification, "dedicated/run-10-attempt-2/job-30", NOW, ) report_claim = store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) self.assertEqual("report", report_claim.stage) self.assertEqual(report, report_claim.report) self.assertEqual(notification, report_claim.notification) store.complete_report(event.event_id, "worker", NOW) notify_claim = store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) self.assertEqual("notify", notify_claim.stage) self.assertEqual(report, notify_claim.report) store.complete_notification(event.event_id, "worker", NOW) self.assertIsNone( store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) ) self.assertEqual("notified", store.get_event(event.event_id).status) def test_failed_notification_retries_same_payload_after_delay(self) -> None: event = failure_event() report = {"event_id": event.event_id, "incidents": []} notification = {"event_id": event.event_id, "incidents": []} with StateStore(self.path) as store: store.discover(event, NOW) store.claim_next( "worker-a", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) store.complete_triage( event.event_id, "worker-a", report, notification, "target/run/job", NOW, ) store.claim_next( "worker-a", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) store.complete_report(event.event_id, "worker-a", NOW) claim = store.claim_next( "worker-a", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) status = store.fail_stage( event.event_id, "worker-a", "timeout token=must-not-leak", NOW, retry_base_seconds=10, retry_max_seconds=60, max_attempts=3, ) self.assertEqual("retry_wait", status) self.assertIsNone( store.claim_next( "worker-b", NOW + timedelta(seconds=9), target_ids=TARGET_IDS, lease_seconds=60, ) ) resumed = store.claim_next( "worker-b", NOW + timedelta(seconds=10), target_ids=TARGET_IDS, lease_seconds=60, ) self.assertEqual("notify", resumed.stage) self.assertEqual(claim.notification, resumed.notification) self.assertNotIn( "must-not-leak", store.get_event(event.event_id).last_error ) def test_retry_limit_moves_event_to_abandoned(self) -> None: event = failure_event() with StateStore(self.path) as store: store.discover(event, NOW) store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) status = store.fail_stage( event.event_id, "worker", "permanent", NOW, retry_base_seconds=1, retry_max_seconds=2, max_attempts=1, ) self.assertEqual("abandoned", status) self.assertIsNone( store.claim_next( "worker", NOW + timedelta(days=1), target_ids=TARGET_IDS, lease_seconds=60, ) ) def test_successful_stage_resets_retry_budget_for_next_stage(self) -> None: event = failure_event() with StateStore(self.path) as store: store.discover(event, NOW) store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) store.fail_stage( event.event_id, "worker", "triage failed once", NOW, retry_base_seconds=1, retry_max_seconds=2, max_attempts=2, ) store.claim_next( "worker", NOW + timedelta(seconds=1), target_ids=TARGET_IDS, lease_seconds=60, ) store.complete_triage( event.event_id, "worker", {}, {}, "event", NOW + timedelta(seconds=1), ) report_claim = store.claim_next( "worker", NOW + timedelta(seconds=1), target_ids=TARGET_IDS, lease_seconds=60, ) self.assertEqual(0, report_claim.attempt_count) status = store.fail_stage( event.event_id, "worker", "report failed once", NOW + timedelta(seconds=1), retry_base_seconds=1, retry_max_seconds=2, max_attempts=2, ) self.assertEqual("retry_wait", status) def test_prune_deletes_only_old_terminal_events(self) -> None: old = failure_event(run_id=1, job_id=1) active = failure_event(run_id=2, job_id=2) with StateStore(self.path) as store: for event in (old, active): store.discover(event, NOW) store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) store.complete_triage(old.event_id, "worker", {}, {}, "old", NOW) store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) store.complete_report(old.event_id, "worker", NOW) store.claim_next( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) store.complete_notification(old.event_id, "worker", NOW) deleted = store.prune(NOW + timedelta(days=31), retention_days=30) self.assertEqual(1, deleted) self.assertIsNone(store.get_event(old.event_id)) self.assertIsNotNone(store.get_event(active.event_id)) def test_notification_batch_only_returns_allowed_targets(self) -> None: other_event = failure_event( target=OTHER_TARGET, run_id=9, job_id=29, ) target_event = failure_event(run_id=10, job_id=30) with StateStore(self.path) as store: for event in (other_event, target_event): store.discover(event, NOW) store.claim_next( "setup", NOW, target_ids=(event.target_id,), lease_seconds=60, ) store.complete_triage( event.event_id, "setup", {"event_id": event.event_id}, {"event_id": event.event_id}, event.event_id, NOW, ) store.claim_next( "setup", NOW, target_ids=(event.target_id,), lease_seconds=60, ) store.complete_report(event.event_id, "setup", NOW) claims = store.claim_notification_batch( "worker", NOW, target_ids=TARGET_IDS, lease_seconds=60, ) other_record = store.get_event(other_event.event_id) assert other_record is not None self.assertEqual( [target_event.event_id], [claim.event.event_id for claim in claims], ) self.assertEqual( "report_ready", other_record.status, ) if __name__ == "__main__": unittest.main()