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 cloud.database import CloudDatabase
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()