agentic-mobile-control/tests/test_cloud_repository_concurrency.py

from __future__ import annotations

import os
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime, timedelta
from pathlib import Path
from threading import Barrier
from uuid import uuid4

import pytest

from cloud.database import CloudDatabase
from cloud.pool import PooledDevice
from cloud.scheduler import ScheduledTask, TaskConstraints


def _unique_id(prefix: str) -> str:
    return f"{prefix}-{uuid4().hex}"


def _seed_assignment_candidate(
    database: CloudDatabase,
    *,
    now: datetime,
) -> tuple[str, str, str]:
    host_id = _unique_id("concurrency-host")
    device_id = _unique_id("concurrency-device")
    task_id = _unique_id("concurrency-task")
    database.repository.upsert_host(host_id, address=None, last_seen_at=now)
    database.repository.replace_host_devices(
        host_id,
        [
            PooledDevice(
                device_id=device_id,
                host_id=host_id,
                driver_type="wda",
                status="idle",
                synced_at=now,
            )
        ],
    )
    database.repository.enqueue_task(
        ScheduledTask(
            id=task_id,
            goal="concurrent operation",
            workflow_definition_id=None,
            constraints=TaskConstraints(),
            created_at=now,
        )
    )
    return host_id, device_id, task_id


def _postgres_url() -> str:
    url = os.getenv("TEST_POSTGRES_URL")
    if not url:
        pytest.skip("TEST_POSTGRES_URL is required for PostgreSQL concurrency tests")
    return url


@pytest.mark.integration
def test_postgresql_concurrent_assignment_has_one_winner() -> None:
    database_url = _postgres_url()
    setup = CloudDatabase(database_url)
    now = datetime(2026, 7, 12, 17, 0, tzinfo=UTC)
    host_id, device_id, task_id = _seed_assignment_candidate(setup, now=now)
    setup.close()
    barrier = Barrier(2)

    def assign(lease_id: str):
        contender = CloudDatabase(database_url, create_schema=False)
        try:
            barrier.wait(timeout=5)
            return contender.repository.assign_task(
                task_id=task_id,
                host_id=host_id,
                device_id=device_id,
                lease_id=lease_id,
                lease_expires_at=now + timedelta(minutes=1),
                now=now,
            )
        finally:
            contender.close()

    with ThreadPoolExecutor(max_workers=2) as executor:
        results = list(executor.map(assign, ("lease-a", "lease-b")))

    assert sum(result is not None for result in results) == 1
    verification = CloudDatabase(database_url, create_schema=False)
    try:
        task = verification.repository.get_task(task_id)
        attempts = verification.repository.list_task_attempts(task_id)
        assert task is not None
        assert task.status == "assigned"
        assert task.attempt_count == 1
        assert len(attempts) == 1
        assert attempts[0].lease_id in {"lease-a", "lease-b"}
    finally:
        verification.close()


@pytest.mark.integration
def test_postgresql_concurrent_claim_has_one_winner() -> None:
    database_url = _postgres_url()
    setup = CloudDatabase(database_url)
    now = datetime(2026, 7, 12, 18, 0, tzinfo=UTC)
    host_id, device_id, task_id = _seed_assignment_candidate(setup, now=now)
    setup.repository.assign_task(
        task_id=task_id,
        host_id=host_id,
        device_id=device_id,
        lease_id="claim-race-lease",
        lease_expires_at=now + timedelta(minutes=1),
        now=now,
    )
    setup.close()
    barrier = Barrier(2)

    def claim():
        contender = CloudDatabase(database_url, create_schema=False)
        try:
            barrier.wait(timeout=5)
            return contender.repository.claim_assignment(host_id=host_id, now=now)
        finally:
            contender.close()

    with ThreadPoolExecutor(max_workers=2) as executor:
        results = list(executor.map(lambda _index: claim(), range(2)))

    assert sum(result is not None for result in results) == 1
    verification = CloudDatabase(database_url, create_schema=False)
    try:
        task = verification.repository.get_task(task_id)
        attempts = verification.repository.list_task_attempts(task_id)
        assert task is not None
        assert task.status == "dispatched"
        assert len(attempts) == 1
        assert attempts[0].status == "dispatched"
    finally:
        verification.close()


def test_sqlite_single_control_plane_rejects_sequential_duplicate_assignment(
    tmp_path: Path,
) -> None:
    database = CloudDatabase(f"sqlite:///{(tmp_path / 'single.sqlite3').as_posix()}")
    now = datetime(2026, 7, 12, 19, 0, tzinfo=UTC)
    host_id, device_id, task_id = _seed_assignment_candidate(database, now=now)

    try:
        first = database.repository.assign_task(
            task_id=task_id,
            host_id=host_id,
            device_id=device_id,
            lease_id="sqlite-lease-a",
            lease_expires_at=now + timedelta(minutes=1),
            now=now,
        )
        second = database.repository.assign_task(
            task_id=task_id,
            host_id=host_id,
            device_id=device_id,
            lease_id="sqlite-lease-b",
            lease_expires_at=now + timedelta(minutes=1),
            now=now,
        )

        assert database.engine.dialect.name == "sqlite"
        assert first is not None
        assert second is None
        assert len(database.repository.list_task_attempts(task_id)) == 1
    finally:
        database.close()