"""Tests for the pre-restart drain proof and hard gate (#661). Covers the acceptance criteria: 1. Restart apply without a proof fails closed. 2. A successful drain produces a verifiable proof. 3. An open unsafe mutation makes the proof fail (multi-session fixture). 4. Pass / fail / expired verification paths. Plus the security posture: forged/tampered proofs are rejected, break-glass is the only bypass and is never silent, a stale blast-radius fingerprint rejects a proof, and no per-process secret ever leaks into a serialized artifact. """ from __future__ import annotations import os import unittest from datetime import datetime, timedelta, timezone import drain_proof as dp import restart_coordinator as rc NOW = datetime(2026, 7, 24, 6, 0, 0, tzinfo=timezone.utc) SECRET = b"unit-test-drain-proof-secret-0123456789abcdef" def _live_pid() -> int: return os.getpid() def _clean_drain_state() -> dict: """Every drain action succeeded, no sessions outstanding.""" return { "assignments_stopped": True, "checkpoints_complete": True, "handoffs_verified": True, "leases_handled": True, "acks": {}, # no other live sessions to acknowledge "ack_timeout_policy_applied": False, } def _safe_report() -> dict: """Impact report with no other live work: a restart here is safe.""" report = rc.evaluate_restart_impact( {"sessions": [], "leases": [], "inventory_complete": True}, now=NOW, requesting_session_id="prgs-controller-1-req", ) return report.as_dict() def _unsafe_mutation_report() -> dict: """Multi-session report: a second session holds a live author mutation.""" sessions = [ { "session_id": "prgs-controller-1-req", "role": "controller", "profile": "prgs-controller", "pid": _live_pid(), "status": "active", "last_heartbeat_at": NOW.isoformat(), }, { "session_id": "prgs-author-99", "role": "author", "profile": "prgs-author", "pid": _live_pid(), "status": "active", "last_heartbeat_at": NOW.isoformat(), }, ] leases = [ { "lease_id": "lease-mut", "session_id": "prgs-author-99", "role": "author", "phase": "implementing", "work_kind": "issue", "work_number": 661, "worktree_path": "branches/issue-661", "freshness": {"freshness": "active"}, } ] report = rc.evaluate_restart_impact( {"sessions": sessions, "leases": leases, "inventory_complete": True}, now=NOW, requesting_session_id="prgs-controller-1-req", ) return report.as_dict() class BuildDrainProofTests(unittest.TestCase): def test_clean_drain_produces_verifiable_clean_proof(self): """AC#2: a successful drain produces a verifiable proof.""" proof = dp.build_drain_proof( impact_report=_safe_report(), drain_state=_clean_drain_state(), requesting_session_id="prgs-controller-1-req", now=NOW, secret=SECRET, ) self.assertTrue(proof.clean) self.assertEqual(proof.failed_checks, []) self.assertEqual( {c.name for c in proof.checks}, set(dp.REQUIRED_CHECKS) ) result = dp.verify_drain_proof( proof.as_dict(), now=NOW, secret=SECRET ) self.assertTrue(result.valid, result.reasons) self.assertFalse(result.expired) self.assertFalse(result.tampered) def test_open_mutation_makes_proof_unclean(self): """AC#3: an unsafe mutation still in flight fails the proof.""" proof = dp.build_drain_proof( impact_report=_unsafe_mutation_report(), drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ) self.assertFalse(proof.clean) self.assertIn(dp.CHECK_NO_INFLIGHT_MUTATIONS, proof.failed_checks) # Leases-handled also fails: the report still shows a disruptive lease. self.assertIn(dp.CHECK_LEASES_HANDLED, proof.failed_checks) result = dp.verify_drain_proof(proof.as_dict(), now=NOW, secret=SECRET) self.assertFalse(result.valid) def test_incomplete_inventory_fails_no_mutations_check(self): proof = dp.build_drain_proof( impact_report={"inventory_complete": False}, drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ) self.assertFalse(proof.clean) self.assertIn(dp.CHECK_NO_INFLIGHT_MUTATIONS, proof.failed_checks) def test_missing_checkpoint_flag_fails_closed(self): state = _clean_drain_state() del state["checkpoints_complete"] proof = dp.build_drain_proof( impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET ) self.assertFalse(proof.clean) self.assertIn(dp.CHECK_CHECKPOINTS_COMPLETE, proof.failed_checks) def test_non_true_flags_fail_closed(self): """A truthy-but-not-True value (e.g. the string 'yes') must not pass.""" state = _clean_drain_state() state["assignments_stopped"] = "yes" proof = dp.build_drain_proof( impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET ) self.assertIn(dp.CHECK_ASSIGNMENTS_STOPPED, proof.failed_checks) def test_ack_timeout_policy_satisfies_ack_check(self): state = _clean_drain_state() state["acks"] = {"prgs-author-99": "pending"} state["ack_timeout_policy_applied"] = True proof = dp.build_drain_proof( impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET ) names = {c.name: c.passed for c in proof.checks} self.assertTrue(names[dp.CHECK_ACKS_OR_TIMEOUT]) def test_outstanding_acks_without_timeout_fail(self): state = _clean_drain_state() state["acks"] = {"prgs-author-99": "pending"} state["ack_timeout_policy_applied"] = False proof = dp.build_drain_proof( impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET ) self.assertIn(dp.CHECK_ACKS_OR_TIMEOUT, proof.failed_checks) def test_all_acked_satisfies_ack_check(self): state = _clean_drain_state() state["acks"] = {"prgs-author-99": "acked", "prgs-author-2": "acknowledged"} proof = dp.build_drain_proof( impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET ) names = {c.name: c.passed for c in proof.checks} self.assertTrue(names[dp.CHECK_ACKS_OR_TIMEOUT]) class VerifyDrainProofTests(unittest.TestCase): def _clean_proof_dict(self) -> dict: return dp.build_drain_proof( impact_report=_safe_report(), drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ).as_dict() def test_missing_proof_is_invalid(self): result = dp.verify_drain_proof(None, now=NOW, secret=SECRET) self.assertFalse(result.valid) self.assertIsNone(result.proof_id) def test_expired_proof_is_invalid(self): """AC#4: an expired proof fails verification.""" proof = self._clean_proof_dict() later = NOW + timedelta(seconds=dp.DEFAULT_PROOF_TTL_SECONDS + 1) result = dp.verify_drain_proof(proof, now=later, secret=SECRET) self.assertFalse(result.valid) self.assertTrue(result.expired) def test_proof_valid_just_before_expiry(self): proof = self._clean_proof_dict() almost = NOW + timedelta(seconds=dp.DEFAULT_PROOF_TTL_SECONDS - 1) result = dp.verify_drain_proof(proof, now=almost, secret=SECRET) self.assertTrue(result.valid, result.reasons) def test_wrong_secret_rejected(self): """A proof minted in a prior process (different secret) will not verify.""" proof = self._clean_proof_dict() result = dp.verify_drain_proof(proof, now=NOW, secret=b"other-secret") self.assertFalse(result.valid) self.assertTrue(result.tampered) def test_flipping_clean_flag_is_detected(self): """Forging clean=True on an unclean proof breaks the signature.""" unclean = dp.build_drain_proof( impact_report=_unsafe_mutation_report(), drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ).as_dict() self.assertFalse(unclean["clean"]) unclean["clean"] = True # forge result = dp.verify_drain_proof(unclean, now=NOW, secret=SECRET) self.assertFalse(result.valid) self.assertTrue(result.tampered) def test_tampering_a_check_is_detected(self): unclean = dp.build_drain_proof( impact_report=_unsafe_mutation_report(), drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ).as_dict() for c in unclean["checks"]: if c["name"] == dp.CHECK_NO_INFLIGHT_MUTATIONS: c["passed"] = True # forge the failing check to pass result = dp.verify_drain_proof(unclean, now=NOW, secret=SECRET) self.assertFalse(result.valid) self.assertTrue(result.tampered) def test_missing_required_check_rejected(self): proof = self._clean_proof_dict() proof["checks"] = [ c for c in proof["checks"] if c["name"] != dp.CHECK_HANDOFFS_OK ] result = dp.verify_drain_proof(proof, now=NOW, secret=SECRET) self.assertFalse(result.valid) def test_stale_fingerprint_rejected(self): proof = self._clean_proof_dict() result = dp.verify_drain_proof( proof, now=NOW, secret=SECRET, expected_impact_fingerprint="deadbeef", ) self.assertFalse(result.valid) def test_matching_fingerprint_accepted(self): report = _safe_report() proof = dp.build_drain_proof( impact_report=report, drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ).as_dict() fp = dp.impact_fingerprint(report) result = dp.verify_drain_proof( proof, now=NOW, secret=SECRET, expected_impact_fingerprint=fp ) self.assertTrue(result.valid, result.reasons) class GateApplyRestartTests(unittest.TestCase): def _clean_proof_dict(self) -> dict: return dp.build_drain_proof( impact_report=_safe_report(), drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ).as_dict() def test_apply_without_proof_denied(self): """AC#1: restart apply without a proof fails closed + raises incident.""" decision = dp.gate_apply_restart(proof=None, now=NOW, secret=SECRET) self.assertFalse(decision.allow) self.assertEqual(decision.verdict, dp.GATE_DENY) self.assertIsNotNone(decision.incident) self.assertEqual( decision.incident["kind"], "restart_drain_gate_denied" ) def test_apply_with_valid_proof_allowed(self): decision = dp.gate_apply_restart( proof=self._clean_proof_dict(), now=NOW, secret=SECRET ) self.assertTrue(decision.allow) self.assertEqual(decision.verdict, dp.GATE_ALLOW) self.assertIsNone(decision.incident) def test_apply_with_expired_proof_denied_with_incident(self): later = NOW + timedelta(seconds=dp.DEFAULT_PROOF_TTL_SECONDS + 5) decision = dp.gate_apply_restart( proof=self._clean_proof_dict(), now=later, secret=SECRET ) self.assertFalse(decision.allow) self.assertIsNotNone(decision.incident) def test_apply_with_unclean_proof_denied(self): """AC#3 at the gate: an unsafe-mutation proof is denied.""" unclean = dp.build_drain_proof( impact_report=_unsafe_mutation_report(), drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ).as_dict() decision = dp.gate_apply_restart(proof=unclean, now=NOW, secret=SECRET) self.assertFalse(decision.allow) self.assertIsNotNone(decision.incident) def test_break_glass_allows_without_proof_but_records_bypass(self): decision = dp.gate_apply_restart( proof=None, now=NOW, secret=SECRET, break_glass=True ) self.assertTrue(decision.allow) self.assertEqual(decision.verdict, dp.GATE_BREAK_GLASS) self.assertTrue(decision.break_glass) self.assertIsNone(decision.incident) self.assertTrue(decision.audit_record["break_glass"]) def test_denied_gate_carries_stale_fingerprint_reason(self): decision = dp.gate_apply_restart( proof=self._clean_proof_dict(), now=NOW, secret=SECRET, expected_impact_fingerprint="not-the-fingerprint", ) self.assertFalse(decision.allow) class SecretHygieneTests(unittest.TestCase): def test_secret_never_serialized(self): proof = dp.build_drain_proof( impact_report=_safe_report(), drain_state=_clean_drain_state(), now=NOW, secret=SECRET, ) blob = dp._canonical(proof.as_dict()) self.assertNotIn(SECRET.decode(), blob) # The signature is a hex digest, not the raw secret. self.assertNotIn(SECRET.hex(), blob) def test_incident_descriptor_has_no_secret(self): decision = dp.gate_apply_restart(proof=None, now=NOW, secret=SECRET) blob = dp._canonical(decision.incident) self.assertNotIn(SECRET.decode(), blob) def _drained_report_with_live_sessions(count: int) -> dict: """Report with ``count`` other live sessions but nothing in flight. Every other checklist item passes against this report, so a failure isolates the acknowledgement check rather than tripping on mutations. """ sessions = [ { "session_id": "prgs-controller-1-req", "role": "controller", "profile": "prgs-controller", "pid": _live_pid(), "status": "active", "last_heartbeat_at": NOW.isoformat(), } ] for index in range(count): sessions.append( { "session_id": f"prgs-author-{index}", "role": "author", "profile": "prgs-author", "pid": _live_pid(), "status": "active", "last_heartbeat_at": NOW.isoformat(), } ) report = rc.evaluate_restart_impact( {"sessions": sessions, "leases": [], "inventory_complete": True}, now=NOW, requesting_session_id="prgs-controller-1-req", ) return report.as_dict() class AcknowledgementFailClosedTests(unittest.TestCase): """Acknowledgement evidence must fail closed unless explicitly verified. Regression cover for the reviewed fail-open on PR #882: an absent ``acks`` key collapsed to ``{}`` and was read as "no other live sessions required to acknowledge", so a proof minted clean and the restart gate allowed while the impact report still showed other live sessions. """ def _state(self, **overrides) -> dict: state = _clean_drain_state() state.pop("acks", None) state["ack_timeout_policy_applied"] = False state.update(overrides) return state def _acks_check(self, proof) -> dp.DrainCheck: return next(c for c in proof.checks if c.name == dp.CHECK_ACKS_OR_TIMEOUT) def _build(self, report: dict, state: dict): return dp.build_drain_proof( impact_report=report, drain_state=state, now=NOW, secret=SECRET ) def assertAcksFailClosed(self, report: dict, state: dict) -> None: proof = self._build(report, state) self.assertFalse(self._acks_check(proof).passed) self.assertIn(dp.CHECK_ACKS_OR_TIMEOUT, proof.failed_checks) self.assertFalse(proof.clean) # --- missing / null / empty / malformed ------------------------------ def test_missing_acks_key_with_live_sessions_fails_closed(self): """The exact reviewed defect: absent key, three other live sessions.""" report = _drained_report_with_live_sessions(3) self.assertEqual(report["counts"]["sessions_live_other"], 3) state = self._state() self.assertNotIn("acks", state) proof = self._build(report, state) check = self._acks_check(proof) self.assertFalse(check.passed) self.assertNotIn("no other live sessions", check.detail) self.assertIn("fail closed", check.detail) self.assertFalse(proof.clean) self.assertEqual(proof.failed_checks, [dp.CHECK_ACKS_OR_TIMEOUT]) def test_none_acks_with_live_sessions_fails_closed(self): self.assertAcksFailClosed( _drained_report_with_live_sessions(2), self._state(acks=None) ) def test_empty_acks_with_live_sessions_fails_closed(self): self.assertAcksFailClosed( _drained_report_with_live_sessions(1), self._state(acks={}) ) def test_malformed_acks_fail_closed(self): for malformed in ([], "ack", 7, ("ack",), True): with self.subTest(malformed=malformed): self.assertAcksFailClosed( _drained_report_with_live_sessions(1), self._state(acks=malformed), ) # --- stale / unproven values ----------------------------------------- def test_stale_or_unproven_ack_values_fail_closed(self): for value in ("pending", "stale", "unknown", "", None, True, 1, NOW): with self.subTest(value=value): self.assertAcksFailClosed( _drained_report_with_live_sessions(1), self._state(acks={"prgs-author-0": value}), ) def test_partial_coverage_fails_closed(self): """Fewer acknowledgements than the report's live-session count.""" self.assertAcksFailClosed( _drained_report_with_live_sessions(3), self._state(acks={"prgs-author-0": "ack"}), ) def test_one_unacked_entry_among_many_fails_closed(self): self.assertAcksFailClosed( _drained_report_with_live_sessions(2), self._state(acks={"prgs-author-0": "ack", "prgs-author-1": "pending"}), ) def test_unproven_live_session_count_fails_closed(self): """A missing or malformed count cannot prove nobody had to acknowledge.""" malformed_counts = ( None, {}, {"sessions_live_other": None}, {"sessions_live_other": "3"}, {"sessions_live_other": -1}, {"sessions_live_other": True}, ) for counts in malformed_counts: with self.subTest(counts=counts): report = _drained_report_with_live_sessions(0) if counts is None: report.pop("counts", None) else: report["counts"] = counts self.assertAcksFailClosed(report, self._state()) # --- valid evidence still passes ------------------------------------- def test_complete_valid_acks_pass(self): report = _drained_report_with_live_sessions(2) state = self._state( acks={"prgs-author-0": "ack", "prgs-author-1": "acknowledged"} ) proof = self._build(report, state) self.assertTrue(self._acks_check(proof).passed) self.assertTrue(proof.clean) self.assertEqual(proof.failed_checks, []) def test_no_other_live_sessions_still_passes(self): """Intended behavior retained: zero live sessions needs no acks.""" report = _drained_report_with_live_sessions(0) self.assertEqual(report["counts"]["sessions_live_other"], 0) proof = self._build(report, self._state()) check = self._acks_check(proof) self.assertTrue(check.passed) self.assertIn("sessions_live_other=0", check.detail) self.assertTrue(proof.clean) # --- timeout policy cannot become a second fail-open ------------------ def test_unproven_timeout_policy_cannot_open_the_gate(self): for value in (None, "true", "yes", 1, "True", [], {}): with self.subTest(value=value): self.assertAcksFailClosed( _drained_report_with_live_sessions(2), self._state(ack_timeout_policy_applied=value), ) def test_explicit_timeout_policy_permits(self): proof = self._build( _drained_report_with_live_sessions(2), self._state(ack_timeout_policy_applied=True), ) check = self._acks_check(proof) self.assertTrue(check.passed) self.assertIn("timeout policy", check.detail) self.assertTrue(proof.clean) # --- the gate itself must deny --------------------------------------- def test_failed_ack_check_denies_the_restart_gate(self): report = _drained_report_with_live_sessions(3) proof = self._build(report, self._state()) self.assertFalse(proof.clean) decision = dp.gate_apply_restart( proof=proof.as_dict(), now=NOW, secret=SECRET, expected_impact_fingerprint=dp.impact_fingerprint(report), ) self.assertFalse(decision.allow) self.assertEqual(decision.verdict, dp.GATE_DENY) self.assertIsNotNone(decision.incident) def test_unclean_ack_proof_fails_verification(self): report = _drained_report_with_live_sessions(3) proof = self._build(report, self._state()) result = dp.verify_drain_proof( proof.as_dict(), now=NOW, secret=SECRET, expected_impact_fingerprint=dp.impact_fingerprint(report), ) self.assertFalse(result.valid) self.assertFalse(result.clean) if __name__ == "__main__": unittest.main()