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()