440 lines
14 KiB
Python
440 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import get_protocol_members
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from sqlalchemy import event
|
|
from sqlalchemy.orm import Session
|
|
|
|
from cloud.database import CloudDatabase
|
|
from cloud.db_models import TaskAttemptRow
|
|
from cloud.plugins import PluginManifest
|
|
from cloud.pool import PooledDevice
|
|
from cloud.repository import CloudRepository, LeasedAssignment, TaskAttemptRecord
|
|
from cloud.scheduler import ScheduledTask, TaskConstraints
|
|
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
pytest.param("sqlite", id="sqlite"),
|
|
pytest.param("postgresql", id="postgresql", marks=pytest.mark.integration),
|
|
]
|
|
)
|
|
def database_url(request: pytest.FixtureRequest, tmp_path: Path) -> str:
|
|
if request.param == "sqlite":
|
|
return f"sqlite:///{(tmp_path / 'contract.sqlite3').as_posix()}"
|
|
|
|
url = os.getenv("TEST_POSTGRES_URL")
|
|
if not url:
|
|
pytest.skip("TEST_POSTGRES_URL is required for PostgreSQL contract tests")
|
|
return url
|
|
|
|
|
|
def _unique_id(prefix: str) -> str:
|
|
return f"{prefix}-{uuid4().hex}"
|
|
|
|
|
|
def _device(device_id: str, host_id: str) -> PooledDevice:
|
|
return PooledDevice(
|
|
device_id=device_id,
|
|
host_id=host_id,
|
|
driver_type="wda",
|
|
status="idle",
|
|
capability_tags=["ios", "physical"],
|
|
synced_at=datetime(2026, 7, 12, tzinfo=UTC),
|
|
)
|
|
|
|
|
|
def test_cloud_repository_exposes_crud_and_atomic_lease_operations() -> None:
|
|
members = get_protocol_members(CloudRepository)
|
|
|
|
assert {
|
|
"upsert_host",
|
|
"replace_host_devices",
|
|
"list_hosts",
|
|
"list_devices",
|
|
"enqueue_task",
|
|
"get_task",
|
|
"save_plugin",
|
|
"assign_task",
|
|
"claim_assignment",
|
|
"renew_lease",
|
|
"record_task_result",
|
|
"reap_expired_leases",
|
|
"list_task_attempts",
|
|
"health_check",
|
|
"close",
|
|
} <= members
|
|
|
|
|
|
def test_repository_transfer_records_are_immutable() -> None:
|
|
assert TaskAttemptRecord.__dataclass_params__.frozen is True
|
|
assert LeasedAssignment.__dataclass_params__.frozen is True
|
|
|
|
|
|
def test_repository_crud_contract(database_url: str) -> None:
|
|
database = CloudDatabase(database_url)
|
|
repository = database.repository
|
|
host_id = _unique_id("host")
|
|
device_id = _unique_id("device")
|
|
task_id = _unique_id("task")
|
|
plugin_name = _unique_id("plugin")
|
|
seen_at = datetime(2026, 7, 12, 1, 2, 3, tzinfo=UTC)
|
|
|
|
try:
|
|
repository.upsert_host(
|
|
host_id,
|
|
address="127.0.0.1:9000",
|
|
last_seen_at=seen_at,
|
|
)
|
|
repository.replace_host_devices(host_id, [_device(device_id, host_id)])
|
|
repository.enqueue_task(
|
|
ScheduledTask(
|
|
id=task_id,
|
|
goal="open settings",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(
|
|
driver_type="wda",
|
|
capability_tags=["ios"],
|
|
),
|
|
created_at=seen_at,
|
|
)
|
|
)
|
|
repository.save_plugin(
|
|
PluginManifest(
|
|
name=plugin_name,
|
|
version="1.0.0",
|
|
entry_point_kind="tool",
|
|
target="cloud.store:CloudStore",
|
|
),
|
|
wired=False,
|
|
)
|
|
|
|
assert repository.get_host(host_id) is not None
|
|
assert repository.get_device(device_id) == _device(device_id, host_id)
|
|
assert repository.get_task(task_id) is not None
|
|
assert repository.get_plugin(plugin_name) is not None
|
|
repository.health_check()
|
|
finally:
|
|
database.close()
|
|
|
|
|
|
def test_failed_snapshot_transaction_rolls_back(database_url: str) -> None:
|
|
database = CloudDatabase(database_url)
|
|
repository = database.repository
|
|
host_id = _unique_id("rollback-host")
|
|
original_device_id = _unique_id("original-device")
|
|
duplicate_device_id = _unique_id("duplicate-device")
|
|
|
|
try:
|
|
repository.upsert_host(
|
|
host_id,
|
|
address=None,
|
|
last_seen_at=datetime(2026, 7, 12, tzinfo=UTC),
|
|
)
|
|
repository.replace_host_devices(
|
|
host_id,
|
|
[_device(original_device_id, host_id)],
|
|
)
|
|
|
|
def fail_before_insert(
|
|
_connection: object,
|
|
_cursor: object,
|
|
statement: str,
|
|
_parameters: object,
|
|
_context: object,
|
|
_executemany: bool,
|
|
) -> None:
|
|
if statement.startswith("INSERT INTO pooled_devices"):
|
|
raise RuntimeError("injected snapshot write failure")
|
|
|
|
event.listen(database.engine, "before_cursor_execute", fail_before_insert)
|
|
try:
|
|
with pytest.raises(RuntimeError, match="injected snapshot write failure"):
|
|
repository.replace_host_devices(
|
|
host_id,
|
|
[_device(duplicate_device_id, host_id)],
|
|
)
|
|
finally:
|
|
event.remove(database.engine, "before_cursor_execute", fail_before_insert)
|
|
|
|
devices = [
|
|
device for device in repository.list_devices() if device.host_id == host_id
|
|
]
|
|
assert [device.device_id for device in devices] == [original_device_id]
|
|
repository.health_check()
|
|
finally:
|
|
database.close()
|
|
|
|
|
|
def test_repository_state_survives_application_restart(database_url: str) -> None:
|
|
host_id = _unique_id("restart-host")
|
|
device_id = _unique_id("restart-device")
|
|
first_process = CloudDatabase(database_url)
|
|
try:
|
|
first_process.repository.upsert_host(
|
|
host_id,
|
|
address="host.internal",
|
|
last_seen_at=datetime(2026, 7, 12, tzinfo=UTC),
|
|
)
|
|
first_process.repository.replace_host_devices(
|
|
host_id,
|
|
[_device(device_id, host_id)],
|
|
)
|
|
finally:
|
|
first_process.close()
|
|
|
|
restarted_process = CloudDatabase(database_url)
|
|
try:
|
|
host = restarted_process.repository.get_host(host_id)
|
|
device = restarted_process.repository.get_device(device_id)
|
|
|
|
assert host is not None
|
|
assert host.address == "host.internal"
|
|
assert device == _device(device_id, host_id)
|
|
finally:
|
|
restarted_process.close()
|
|
|
|
|
|
def test_task_lease_and_terminal_fields_round_trip(database_url: str) -> None:
|
|
database = CloudDatabase(database_url)
|
|
task_id = _unique_id("lease-task")
|
|
created_at = datetime(2026, 7, 12, 1, 0, tzinfo=UTC)
|
|
lease_expires_at = datetime(2026, 7, 12, 1, 5, tzinfo=UTC)
|
|
updated_at = datetime(2026, 7, 12, 1, 1, tzinfo=UTC)
|
|
task = ScheduledTask(
|
|
id=task_id,
|
|
goal="capture diagnostics",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
status="failed",
|
|
assigned_device_id="device-a",
|
|
assigned_host_id="host-a",
|
|
attempt_count=2,
|
|
lease_id="lease-a",
|
|
lease_expires_at=lease_expires_at,
|
|
terminal_result={"steps": 3, "status": "failed"},
|
|
failure_reason="device disconnected",
|
|
updated_at=updated_at,
|
|
created_at=created_at,
|
|
)
|
|
|
|
try:
|
|
database.repository.enqueue_task(task)
|
|
|
|
assert database.repository.get_task(task_id) == task
|
|
finally:
|
|
database.close()
|
|
|
|
|
|
def test_task_attempt_history_is_ordered_and_complete(database_url: str) -> None:
|
|
database = CloudDatabase(database_url)
|
|
task_id = _unique_id("attempt-task")
|
|
created_at = datetime(2026, 7, 12, 1, 0, tzinfo=UTC)
|
|
lease_expires_at = datetime(2026, 7, 12, 1, 5, tzinfo=UTC)
|
|
|
|
try:
|
|
with Session(database.engine) as session, session.begin():
|
|
session.add_all(
|
|
[
|
|
TaskAttemptRow(
|
|
task_id=task_id,
|
|
attempt=2,
|
|
lease_id="lease-2",
|
|
host_id="host-b",
|
|
device_id="device-b",
|
|
status="failed",
|
|
lease_expires_at=lease_expires_at.isoformat(),
|
|
created_at=created_at.isoformat(),
|
|
completed_at=lease_expires_at.isoformat(),
|
|
failure_reason="execution failed",
|
|
result_json='{"exit_code": 1}',
|
|
),
|
|
TaskAttemptRow(
|
|
task_id=task_id,
|
|
attempt=1,
|
|
lease_id="lease-1",
|
|
host_id="host-a",
|
|
device_id="device-a",
|
|
status="expired",
|
|
lease_expires_at=lease_expires_at.isoformat(),
|
|
created_at=created_at.isoformat(),
|
|
completed_at=lease_expires_at.isoformat(),
|
|
failure_reason="lease expired",
|
|
result_json=None,
|
|
),
|
|
]
|
|
)
|
|
|
|
attempts = database.repository.list_task_attempts(task_id)
|
|
|
|
assert [attempt.attempt for attempt in attempts] == [1, 2]
|
|
assert attempts[0].status == "expired"
|
|
assert attempts[1].terminal_result == {"exit_code": 1}
|
|
assert attempts[1].failure_reason == "execution failed"
|
|
finally:
|
|
database.close()
|
|
|
|
|
|
def test_atomic_assignment_creates_lease_attempt_and_reservation(
|
|
database_url: str,
|
|
) -> None:
|
|
database = CloudDatabase(database_url)
|
|
host_id = _unique_id("assignment-host")
|
|
device_id = _unique_id("assignment-device")
|
|
task_id = _unique_id("assignment-task")
|
|
now = datetime(2026, 7, 12, 2, 0, tzinfo=UTC)
|
|
lease_expires_at = now + timedelta(minutes=1)
|
|
|
|
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="open settings",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
created_at=now,
|
|
)
|
|
)
|
|
|
|
assignment = database.repository.assign_task(
|
|
task_id=task_id,
|
|
host_id=host_id,
|
|
device_id=device_id,
|
|
lease_id="lease-1",
|
|
lease_expires_at=lease_expires_at,
|
|
now=now,
|
|
)
|
|
|
|
assert assignment is not None
|
|
assert assignment.attempt == 1
|
|
assert assignment.lease_id == "lease-1"
|
|
task = database.repository.get_task(task_id)
|
|
assert task is not None
|
|
assert task.status == "assigned"
|
|
assert task.attempt_count == 1
|
|
assert task.lease_expires_at == lease_expires_at
|
|
assert device_id in database.repository.list_reserved_device_ids(now=now)
|
|
attempts = database.repository.list_task_attempts(task_id)
|
|
assert len(attempts) == 1
|
|
assert attempts[0].status == "assigned"
|
|
assert attempts[0].lease_id == "lease-1"
|
|
finally:
|
|
database.close()
|
|
|
|
|
|
def test_active_assignment_blocks_reuse_of_stale_idle_snapshot(
|
|
database_url: str,
|
|
) -> None:
|
|
database = CloudDatabase(database_url)
|
|
host_id = _unique_id("reservation-host")
|
|
device_id = _unique_id("reservation-device")
|
|
first_task_id = _unique_id("reservation-task")
|
|
second_task_id = _unique_id("reservation-task")
|
|
now = datetime(2026, 7, 12, 3, 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)],
|
|
)
|
|
for task_id in (first_task_id, second_task_id):
|
|
database.repository.enqueue_task(
|
|
ScheduledTask(
|
|
id=task_id,
|
|
goal="run task",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
created_at=now,
|
|
)
|
|
)
|
|
|
|
first = database.repository.assign_task(
|
|
task_id=first_task_id,
|
|
host_id=host_id,
|
|
device_id=device_id,
|
|
lease_id="lease-active",
|
|
lease_expires_at=now + timedelta(minutes=1),
|
|
now=now,
|
|
)
|
|
second = database.repository.assign_task(
|
|
task_id=second_task_id,
|
|
host_id=host_id,
|
|
device_id=device_id,
|
|
lease_id="lease-blocked",
|
|
lease_expires_at=now + timedelta(minutes=1),
|
|
now=now,
|
|
)
|
|
|
|
assert first is not None
|
|
assert second is None
|
|
blocked_task = database.repository.get_task(second_task_id)
|
|
assert blocked_task is not None
|
|
assert blocked_task.status == "queued"
|
|
assert database.repository.list_task_attempts(second_task_id) == []
|
|
finally:
|
|
database.close()
|
|
|
|
|
|
def test_expired_assignment_no_longer_reserves_device(database_url: str) -> None:
|
|
database = CloudDatabase(database_url)
|
|
host_id = _unique_id("expired-host")
|
|
device_id = _unique_id("expired-device")
|
|
first_task_id = _unique_id("expired-task")
|
|
second_task_id = _unique_id("expired-task")
|
|
assigned_at = datetime(2026, 7, 12, 4, 0, tzinfo=UTC)
|
|
after_expiry = assigned_at + timedelta(minutes=2)
|
|
|
|
try:
|
|
database.repository.upsert_host(
|
|
host_id,
|
|
address=None,
|
|
last_seen_at=assigned_at,
|
|
)
|
|
database.repository.replace_host_devices(
|
|
host_id,
|
|
[_device(device_id, host_id)],
|
|
)
|
|
for task_id in (first_task_id, second_task_id):
|
|
database.repository.enqueue_task(
|
|
ScheduledTask(
|
|
id=task_id,
|
|
goal="run task",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
created_at=assigned_at,
|
|
)
|
|
)
|
|
assert database.repository.assign_task(
|
|
task_id=first_task_id,
|
|
host_id=host_id,
|
|
device_id=device_id,
|
|
lease_id="lease-expired",
|
|
lease_expires_at=assigned_at + timedelta(minutes=1),
|
|
now=assigned_at,
|
|
)
|
|
|
|
assert device_id not in database.repository.list_reserved_device_ids(
|
|
now=after_expiry
|
|
)
|
|
assert database.repository.assign_task(
|
|
task_id=second_task_id,
|
|
host_id=host_id,
|
|
device_id=device_id,
|
|
lease_id="lease-new",
|
|
lease_expires_at=after_expiry + timedelta(minutes=1),
|
|
now=after_expiry,
|
|
)
|
|
finally:
|
|
database.close()
|