feat(host-protocol): renew and complete leases

2026-07-12 18:15:04 +08:00
parent c7f8d0c128
commit a24efc1043
4 changed files with 242 additions and 1 deletions
@@ -204,3 +204,136 @@ def test_empty_long_poll_timeout_is_normal_response(tmp_path) -> None:

    assert response.status_code == 200
    assert response.json() == {"assignment": None, "timed_out": True}
+
+
+def _seed_active_assignment(pool: DevicePool) -> datetime:
+    now = datetime.now(UTC)
+    pool.store.upsert_host("host-a", address=None, last_seen_at=now)
+    pool.store.replace_host_devices(
+        "host-a",
+        [
+            PooledDevice(
+                device_id="active-device",
+                host_id="host-a",
+                driver_type="wda",
+                status="idle",
+                synced_at=now,
+            )
+        ],
+    )
+    pool.store.enqueue_task(
+        ScheduledTask(
+            id="active-task",
+            goal="execute assignment",
+            workflow_definition_id=None,
+            constraints=TaskConstraints(),
+            created_at=now,
+        )
+    )
+    pool.store.assign_task(
+        task_id="active-task",
+        host_id="host-a",
+        device_id="active-device",
+        lease_id="active-lease",
+        lease_expires_at=now + timedelta(minutes=1),
+        now=now,
+    )
+    pool.store.claim_assignment(host_id="host-a", now=now)
+    return now
+
+
+def test_lease_renewal_extends_active_assignment(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    original_time = _seed_active_assignment(pool)
+
+    response = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/renew",
+        headers={"Authorization": "Bearer token-a"},
+        json={
+            "host_id": "host-a",
+            "task_id": "active-task",
+            "attempt": 1,
+            "lease_id": "active-lease",
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "renewed"
+    renewed_expiry = datetime.fromisoformat(response.json()["lease_expires_at"])
+    assert renewed_expiry > original_time + timedelta(seconds=30)
+
+
+def test_stale_renewal_returns_typed_conflict(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    _seed_active_assignment(pool)
+
+    response = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/renew",
+        headers={"Authorization": "Bearer token-a"},
+        json={
+            "host_id": "host-a",
+            "task_id": "active-task",
+            "attempt": 1,
+            "lease_id": "stale-lease",
+        },
+    )
+
+    assert response.status_code == 409
+    assert response.json()["code"] == "stale_lease"
+
+
+def test_terminal_result_is_idempotent_through_internal_api(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    _seed_active_assignment(pool)
+    payload = {
+        "host_id": "host-a",
+        "task_id": "active-task",
+        "attempt": 1,
+        "lease_id": "active-lease",
+        "status": "done",
+        "result": {"steps": 4},
+    }
+
+    first = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json=payload,
+    )
+    repeated = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json=payload,
+    )
+
+    assert first.status_code == 200
+    assert first.json()["status"] == "recorded"
+    assert repeated.status_code == 200
+    assert repeated.json()["status"] == "already_recorded"
+    assert pool.store.get_task("active-task").status == "done"  # type: ignore[union-attr]
+
+
+def test_conflicting_repeated_result_returns_stale_lease_conflict(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    _seed_active_assignment(pool)
+    base_payload = {
+        "host_id": "host-a",
+        "task_id": "active-task",
+        "attempt": 1,
+        "lease_id": "active-lease",
+        "status": "done",
+        "result": {"steps": 4},
+    }
+    client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json=base_payload,
+    )
+
+    response = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json={**base_payload, "status": "failed", "failure_reason": "late failure"},
+    )
+
+    assert response.status_code == 409
+    assert response.json()["code"] == "stale_lease"