This commit is contained in:
@@ -10,7 +10,7 @@ from cloud.auth import BearerCredential, ConfiguredBearerAuthProvider
|
||||
from cloud.config import CloudConfig
|
||||
from cloud.internal_api.api import create_internal_router
|
||||
from cloud.pool import DevicePool, PooledDevice
|
||||
from cloud.scheduler import ScheduledTask, TaskConstraints
|
||||
from cloud.scheduler import ScheduledTask, TaskConstraints, TaskScheduler
|
||||
from cloud.store import CloudStore
|
||||
|
||||
|
||||
@@ -34,7 +34,13 @@ def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
|
||||
]
|
||||
)
|
||||
app = FastAPI()
|
||||
app.include_router(create_internal_router(pool=pool, auth_provider=auth_provider))
|
||||
app.include_router(
|
||||
create_internal_router(
|
||||
pool=pool,
|
||||
auth_provider=auth_provider,
|
||||
scheduler=TaskScheduler(pool, pool.store, CloudConfig(stale_after_seconds=60)),
|
||||
)
|
||||
)
|
||||
return TestClient(app), pool
|
||||
|
||||
|
||||
@@ -196,6 +202,36 @@ def test_host_token_cannot_submit_heartbeat_for_another_host(tmp_path) -> None:
|
||||
assert pool.store.get_host("host-b") is None
|
||||
|
||||
|
||||
def test_heartbeat_and_self_submission_preserve_host_isolation(tmp_path) -> None:
|
||||
client, pool = _build_client(tmp_path)
|
||||
headers = {"Authorization": "Bearer token-a"}
|
||||
heartbeat = client.put(
|
||||
"/internal/v1/hosts/host-a/heartbeat",
|
||||
headers=headers,
|
||||
json=_heartbeat_payload("host-a", "device-a"),
|
||||
)
|
||||
assert heartbeat.status_code == 200
|
||||
assert heartbeat.json()["policy_revision"] == 0
|
||||
assert heartbeat.json()["policy"] is None
|
||||
|
||||
created = client.post(
|
||||
"/internal/v1/hosts/host-a/tasks",
|
||||
headers=headers,
|
||||
json={"host_id": "host-a", "goal": "local work", "device_id": "device-a"},
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
task = pool.store.get_task(created.json()["task_id"])
|
||||
assert task is not None
|
||||
assert task.constraints.target_host_id == "host-a"
|
||||
assert task.constraints.target_device_id == "device-a"
|
||||
foreign = client.post(
|
||||
"/internal/v1/hosts/host-b/tasks",
|
||||
headers=headers,
|
||||
json={"host_id": "host-b", "goal": "forbidden"},
|
||||
)
|
||||
assert foreign.status_code == 403
|
||||
|
||||
|
||||
def test_long_poll_claim_returns_at_most_one_owned_assignment(tmp_path) -> None:
|
||||
client, pool = _build_client(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
|
||||
Reference in New Issue
Block a user