Pulse/scripts/release_control/work_claim_test.py

from __future__ import annotations

from datetime import datetime, timezone
import unittest

from work_claim import claim_work, parse_args, release_claims


def base_payload() -> dict[str, object]:
    return {
        "lanes": [
            {"id": "L1", "repo_ids": ["pulse"]},
            {"id": "L2", "repo_ids": ["pulse-mobile"]},
        ],
        "lane_followups": [
            {"id": "followup-l1", "lane_ids": ["L1"], "repo_ids": ["pulse"]},
        ],
        "coverage_gaps": [
            {"id": "gap-l1", "lane_ids": ["L1"], "repo_ids": ["pulse"]},
        ],
        "candidate_lanes": [
            {
                "id": "candidate-l1",
                "target_id": "v6-product-lane-expansion",
                "current_lane_ids": ["L1"],
                "coverage_gap_ids": ["gap-l1"],
                "repo_ids": ["pulse"],
            }
        ],
        "readiness_assertions": [
            {"id": "RA1", "lane_ids": ["L1"], "repo_ids": ["pulse"]},
        ],
        "release_gates": [
            {"id": "g1", "lane_ids": ["L1"], "repo_ids": ["pulse"]},
        ],
        "open_decisions": [
            {"id": "d1", "lane_ids": ["L1"], "repo_ids": ["pulse"]},
        ],
        "work_claims": [],
    }


class WorkClaimTest(unittest.TestCase):
    def test_parse_args_accepts_claim_flags(self) -> None:
        args = parse_args(
            [
                "--kind",
                "candidate-lane",
                "--id",
                "candidate-l1",
                "--summary",
                "Advance the candidate lane.",
                "--agent-id",
                "codex",
                "--pretty",
            ]
        )
        self.assertEqual(args.kind, "candidate-lane")
        self.assertEqual(args.work_id, "candidate-l1")
        self.assertEqual(args.agent_id, "codex")
        self.assertTrue(args.pretty)

    def test_claim_work_replaces_same_agent_claim_and_prunes_expired(self) -> None:
        payload = base_payload()
        payload["work_claims"] = [
            {
                "id": "old-active",
                "agent_id": "codex",
                "summary": "Old lane claim.",
                "target_id": "v6-product-lane-expansion",
                "claimed_at": "2026-03-19T09:00:00Z",
                "heartbeat_at": "2026-03-19T09:15:00Z",
                "expires_at": "2026-03-19T11:00:00Z",
                "work_item": {"kind": "lane", "id": "L2"},
            },
            {
                "id": "expired-claim",
                "agent_id": "other-agent",
                "summary": "Expired gap claim.",
                "target_id": "v6-product-lane-expansion",
                "claimed_at": "2026-03-19T06:00:00Z",
                "heartbeat_at": "2026-03-19T06:15:00Z",
                "expires_at": "2026-03-19T07:00:00Z",
                "work_item": {"kind": "coverage-gap", "id": "gap-l1"},
            },
        ]

        result = claim_work(
            payload,
            work_kind="candidate-lane",
            work_id="candidate-l1",
            summary="Advance the candidate lane.",
            agent_id="codex",
            ttl_hours=2,
            now_utc=datetime(2026, 3, 19, 10, 0, tzinfo=timezone.utc),
        )

        self.assertEqual(result["action"], "claimed")
        self.assertEqual(result["claim"]["target_id"], "v6-product-lane-expansion")
        self.assertEqual(result["claim"]["work_item"], {"kind": "candidate-lane", "id": "candidate-l1"})
        self.assertEqual([claim["id"] for claim in result["replaced_claims"]], ["old-active"])
        self.assertEqual([claim["id"] for claim in result["pruned_expired_claims"]], ["expired-claim"])
        self.assertEqual(result["payload"]["work_claims"], [result["claim"]])

    def test_claim_work_renews_exact_same_claim(self) -> None:
        payload = base_payload()
        payload["work_claims"] = [
            {
                "id": "codex-candidate-lane-candidate-l1",
                "agent_id": "codex",
                "summary": "Old summary.",
                "target_id": "v6-product-lane-expansion",
                "claimed_at": "2026-03-19T09:00:00Z",
                "heartbeat_at": "2026-03-19T09:15:00Z",
                "expires_at": "2026-03-19T11:00:00Z",
                "work_item": {"kind": "candidate-lane", "id": "candidate-l1"},
            }
        ]

        result = claim_work(
            payload,
            work_kind="candidate-lane",
            work_id="candidate-l1",
            summary="Renewed summary.",
            agent_id="codex",
            ttl_hours=2,
            now_utc=datetime(2026, 3, 19, 10, 30, tzinfo=timezone.utc),
        )

        self.assertEqual(result["action"], "renewed")
        self.assertEqual(result["claim"]["id"], "codex-candidate-lane-candidate-l1")
        self.assertEqual(result["claim"]["claimed_at"], "2026-03-19T09:00:00Z")
        self.assertEqual(result["claim"]["heartbeat_at"], "2026-03-19T10:30:00Z")
        self.assertEqual(result["claim"]["expires_at"], "2026-03-19T12:30:00Z")
        self.assertEqual(result["claim"]["summary"], "Renewed summary.")

    def test_claim_work_rejects_existing_same_work_item_claim(self) -> None:
        payload = base_payload()
        payload["work_claims"] = [
            {
                "id": "other-agent-claim",
                "agent_id": "other-agent",
                "summary": "Already taking the candidate lane.",
                "target_id": "v6-product-lane-expansion",
                "claimed_at": "2026-03-19T09:00:00Z",
                "heartbeat_at": "2026-03-19T09:15:00Z",
                "expires_at": "2026-03-19T11:00:00Z",
                "work_item": {"kind": "candidate-lane", "id": "candidate-l1"},
            }
        ]

        with self.assertRaisesRegex(ValueError, "already reserves candidate-lane:candidate-l1"):
            claim_work(
                payload,
                work_kind="candidate-lane",
                work_id="candidate-l1",
                summary="Try to steal the candidate lane.",
                agent_id="codex",
                ttl_hours=2,
                now_utc=datetime(2026, 3, 19, 10, 0, tzinfo=timezone.utc),
            )

    def test_release_claims_removes_matching_agent_claim(self) -> None:
        payload = base_payload()
        payload["work_claims"] = [
            {
                "id": "codex-lane-L1",
                "agent_id": "codex",
                "summary": "Advance lane L1.",
                "target_id": "v6-product-lane-expansion",
                "claimed_at": "2026-03-19T09:00:00Z",
                "heartbeat_at": "2026-03-19T09:15:00Z",
                "expires_at": "2026-03-19T11:00:00Z",
                "work_item": {"kind": "lane", "id": "L1"},
            },
            {
                "id": "other-agent-lane-L2",
                "agent_id": "other-agent",
                "summary": "Advance lane L2.",
                "target_id": "v6-product-lane-expansion",
                "claimed_at": "2026-03-19T09:05:00Z",
                "heartbeat_at": "2026-03-19T09:15:00Z",
                "expires_at": "2026-03-19T11:00:00Z",
                "work_item": {"kind": "lane", "id": "L2"},
            },
        ]

        result = release_claims(
            payload,
            agent_id="codex",
            work_kind="lane",
            work_id="L1",
            now_utc=datetime(2026, 3, 19, 10, 0, tzinfo=timezone.utc),
        )

        self.assertEqual(result["action"], "released")
        self.assertEqual([claim["id"] for claim in result["released_claims"]], ["codex-lane-L1"])
        self.assertEqual(
            [claim["id"] for claim in result["payload"]["work_claims"]],
            ["other-agent-lane-L2"],
        )


if __name__ == "__main__":
    unittest.main()