test(host-protocol): verify internal API isolation

2026-07-12 18:17:39 +08:00
parent a24efc1043
commit 09cfe54dd9
2 changed files with 121 additions and 1 deletions
@@ -337,3 +337,123 @@ def test_conflicting_repeated_result_returns_stale_lease_conflict(tmp_path) -> N

    assert response.status_code == 409
    assert response.json()["code"] == "stale_lease"
+
+
+def test_multi_host_claims_are_isolated(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    now = datetime.now(UTC)
+    for host_id, device_id, task_id, lease_id in (
+        ("host-a", "device-a", "task-a", "lease-a"),
+        ("host-b", "device-b", "task-b", "lease-b"),
+    ):
+        pool.store.upsert_host(host_id, address=None, last_seen_at=now)
+        pool.store.replace_host_devices(
+            host_id,
+            [
+                PooledDevice(
+                    device_id=device_id,
+                    host_id=host_id,
+                    driver_type="wda",
+                    status="idle",
+                    synced_at=now,
+                )
+            ],
+        )
+        pool.store.enqueue_task(
+            ScheduledTask(
+                id=task_id,
+                goal=f"work for {host_id}",
+                workflow_definition_id=None,
+                constraints=TaskConstraints(),
+                created_at=now,
+            )
+        )
+        pool.store.assign_task(
+            task_id=task_id,
+            host_id=host_id,
+            device_id=device_id,
+            lease_id=lease_id,
+            lease_expires_at=now + timedelta(minutes=1),
+            now=now,
+        )
+
+    host_a = client.post(
+        "/internal/v1/hosts/host-a/assignments/claim",
+        headers={"Authorization": "Bearer token-a"},
+        json={"host_id": "host-a", "timeout_seconds": 0},
+    )
+    host_b = client.post(
+        "/internal/v1/hosts/host-b/assignments/claim",
+        headers={"Authorization": "Bearer token-b"},
+        json={"host_id": "host-b", "timeout_seconds": 0},
+    )
+
+    assert host_a.json()["assignment"]["task_id"] == "task-a"
+    assert host_b.json()["assignment"]["task_id"] == "task-b"
+
+
+def test_superseded_attempt_result_cannot_overwrite_current_lease(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    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="retry-device",
+                host_id="host-a",
+                driver_type="wda",
+                status="idle",
+                synced_at=now,
+            )
+        ],
+    )
+    pool.store.enqueue_task(
+        ScheduledTask(
+            id="retry-task",
+            goal="retry safely",
+            workflow_definition_id=None,
+            constraints=TaskConstraints(),
+            created_at=now,
+        )
+    )
+    pool.store.assign_task(
+        task_id="retry-task",
+        host_id="host-a",
+        device_id="retry-device",
+        lease_id="old-lease",
+        lease_expires_at=now + timedelta(seconds=1),
+        now=now,
+    )
+    reaped_at = now + timedelta(seconds=2)
+    pool.store.reap_expired_leases(now=reaped_at, max_attempts=2)
+    pool.store.assign_task(
+        task_id="retry-task",
+        host_id="host-a",
+        device_id="retry-device",
+        lease_id="current-lease",
+        lease_expires_at=reaped_at + timedelta(minutes=1),
+        now=reaped_at,
+    )
+
+    response = client.post(
+        "/internal/v1/hosts/host-a/assignments/retry-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json={
+            "host_id": "host-a",
+            "task_id": "retry-task",
+            "attempt": 1,
+            "lease_id": "old-lease",
+            "status": "done",
+            "result": {"late": True},
+        },
+    )
+
+    assert response.status_code == 409
+    assert response.json()["code"] == "stale_lease"
+    task = pool.store.get_task("retry-task")
+    assert task is not None
+    assert task.status == "assigned"
+    assert task.attempt_count == 2
+    assert task.lease_id == "current-lease"
+    assert task.terminal_result is None