feat(cloud-scheduler): claim host assignments

2026-07-12 17:20:51 +08:00
parent ac7734c7dc
commit b0a5b1426f
3 changed files with 184 additions and 1 deletions
@@ -437,3 +437,143 @@ def test_expired_assignment_no_longer_reserves_device(database_url: str) -> None
        )
    finally:
        database.close()
+
+
+def test_owning_host_claims_one_active_assignment(database_url: str) -> None:
+    database = CloudDatabase(database_url)
+    host_id = _unique_id("claim-host")
+    device_id = _unique_id("claim-device")
+    task_id = _unique_id("claim-task")
+    now = datetime(2026, 7, 12, 5, 0, tzinfo=UTC)
+
+    try:
+        database.repository.upsert_host(host_id, address=None, last_seen_at=now)
+        database.repository.replace_host_devices(
+            host_id,
+            [_device(device_id, host_id)],
+        )
+        database.repository.enqueue_task(
+            ScheduledTask(
+                id=task_id,
+                goal="claim me",
+                workflow_definition_id=None,
+                constraints=TaskConstraints(),
+                created_at=now,
+            )
+        )
+        database.repository.assign_task(
+            task_id=task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="claim-lease",
+            lease_expires_at=now + timedelta(minutes=1),
+            now=now,
+        )
+
+        assignment = database.repository.claim_assignment(
+            host_id=host_id,
+            now=now + timedelta(seconds=1),
+        )
+
+        assert assignment is not None
+        assert assignment.task_id == task_id
+        assert assignment.lease_id == "claim-lease"
+        task = database.repository.get_task(task_id)
+        assert task is not None
+        assert task.status == "dispatched"
+        attempts = database.repository.list_task_attempts(task_id)
+        assert len(attempts) == 1
+        assert attempts[0].status == "dispatched"
+        assert database.repository.claim_assignment(host_id=host_id, now=now) is None
+    finally:
+        database.close()
+
+
+def test_foreign_host_cannot_claim_assignment(database_url: str) -> None:
+    database = CloudDatabase(database_url)
+    host_id = _unique_id("owner-host")
+    device_id = _unique_id("owner-device")
+    task_id = _unique_id("owner-task")
+    now = datetime(2026, 7, 12, 6, 0, tzinfo=UTC)
+
+    try:
+        database.repository.upsert_host(host_id, address=None, last_seen_at=now)
+        database.repository.replace_host_devices(
+            host_id,
+            [_device(device_id, host_id)],
+        )
+        database.repository.enqueue_task(
+            ScheduledTask(
+                id=task_id,
+                goal="owner only",
+                workflow_definition_id=None,
+                constraints=TaskConstraints(),
+                created_at=now,
+            )
+        )
+        database.repository.assign_task(
+            task_id=task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="owner-lease",
+            lease_expires_at=now + timedelta(minutes=1),
+            now=now,
+        )
+
+        assert (
+            database.repository.claim_assignment(
+                host_id=_unique_id("foreign-host"),
+                now=now,
+            )
+            is None
+        )
+        task = database.repository.get_task(task_id)
+        assert task is not None
+        assert task.status == "assigned"
+    finally:
+        database.close()
+
+
+def test_expired_assignment_cannot_be_claimed(database_url: str) -> None:
+    database = CloudDatabase(database_url)
+    host_id = _unique_id("expired-claim-host")
+    device_id = _unique_id("expired-claim-device")
+    task_id = _unique_id("expired-claim-task")
+    now = datetime(2026, 7, 12, 7, 0, tzinfo=UTC)
+
+    try:
+        database.repository.upsert_host(host_id, address=None, last_seen_at=now)
+        database.repository.replace_host_devices(
+            host_id,
+            [_device(device_id, host_id)],
+        )
+        database.repository.enqueue_task(
+            ScheduledTask(
+                id=task_id,
+                goal="too late",
+                workflow_definition_id=None,
+                constraints=TaskConstraints(),
+                created_at=now,
+            )
+        )
+        database.repository.assign_task(
+            task_id=task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id="expired-claim-lease",
+            lease_expires_at=now + timedelta(seconds=1),
+            now=now,
+        )
+
+        assert (
+            database.repository.claim_assignment(
+                host_id=host_id,
+                now=now + timedelta(seconds=2),
+            )
+            is None
+        )
+        task = database.repository.get_task(task_id)
+        assert task is not None
+        assert task.status == "assigned"
+    finally:
+        database.close()