agentic-mobile-control/tests/test_cloud_repository_contract.py

from __future__ import annotations

import os
from datetime import UTC, datetime
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()