agentic-mobile-control/tests/test_cloud_store.py

"""Unit tests for cloud.store.CloudStore (task 2.5)."""

from __future__ import annotations

import pytest

from cloud.pool import HostRegistration, PooledDevice
from cloud.store import CloudStore


def _pooled(
    device_id: str,
    host_id: str,
    *,
    status: str = "idle",
    driver_type: str = "wda",
    tags: list[str] | None = None,
) -> PooledDevice:
    from datetime import UTC, datetime

    return PooledDevice(
        device_id=device_id,
        host_id=host_id,
        driver_type=driver_type,
        status=status,  # type: ignore[arg-type]
        capability_tags=list(tags or []),
        synced_at=datetime(2026, 1, 1, tzinfo=UTC),
    )


def test_upsert_host_and_replace_devices_round_trip(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    from datetime import UTC, datetime

    ts = datetime(2026, 7, 6, 12, 0, tzinfo=UTC)
    store.upsert_host("host-a", address="10.0.0.1:8000", last_seen_at=ts)
    store.replace_host_devices(
        "host-a",
        [_pooled("dev-1", "host-a"), _pooled("dev-2", "host-a", status="busy")],
    )

    hosts = store.list_hosts()
    assert len(hosts) == 1
    assert hosts[0].host_id == "host-a"
    assert hosts[0].address == "10.0.0.1:8000"
    assert hosts[0].last_seen_at == ts

    devices = store.list_devices()
    assert {d.device_id for d in devices} == {"dev-1", "dev-2"}
    by_id = {d.device_id: d for d in devices}
    assert by_id["dev-1"].host_id == "host-a"
    assert by_id["dev-2"].status == "busy"

    fetched = store.get_device("dev-1")
    assert fetched is not None
    assert fetched.host_id == "host-a"
    assert store.get_device("does-not-exist") is None


def test_second_sync_fully_replaces_host_devices(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    from datetime import UTC, datetime

    store.upsert_host("host-a", address=None, last_seen_at=datetime(2026, 1, 1, tzinfo=UTC))
    store.replace_host_devices(
        "host-a",
        [_pooled("dev-1", "host-a"), _pooled("dev-2", "host-a"), _pooled("dev-3", "host-a")],
    )

    # Second sync: only dev-2 plus a new dev-4. dev-1/dev-3 must be gone.
    store.replace_host_devices(
        "host-a",
        [_pooled("dev-2", "host-a"), _pooled("dev-4", "host-a")],
    )

    devices = store.list_devices()
    assert {d.device_id for d in devices} == {"dev-2", "dev-4"}
    assert all(d.host_id == "host-a" for d in devices)


def test_devices_from_two_hosts_coexist(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    from datetime import UTC, datetime

    store.upsert_host("host-a", address="a", last_seen_at=datetime(2026, 1, 1, tzinfo=UTC))
    store.upsert_host("host-b", address="b", last_seen_at=datetime(2026, 1, 2, tzinfo=UTC))
    store.replace_host_devices("host-a", [_pooled("a-dev-1", "host-a")])
    store.replace_host_devices("host-b", [_pooled("b-dev-1", "host-b"), _pooled("b-dev-2", "host-b")])

    devices = store.list_devices()
    assert {d.device_id for d in devices} == {"a-dev-1", "b-dev-1", "b-dev-2"}
    by_host = {d.device_id: d.host_id for d in devices}
    assert by_host == {"a-dev-1": "host-a", "b-dev-1": "host-b", "b-dev-2": "host-b"}

    # Replacing host-a's devices must not touch host-b.
    store.replace_host_devices("host-a", [_pooled("a-dev-9", "host-a")])
    devices = store.list_devices()
    assert {d.device_id for d in devices} == {"a-dev-9", "b-dev-1", "b-dev-2"}


def test_upsert_host_preserves_address_when_none(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    from datetime import UTC, datetime

    store.upsert_host("host-a", address="10.0.0.1:8000", last_seen_at=datetime(2026, 1, 1, tzinfo=UTC))
    # Subsequent sync with address=None should not clobber the existing address.
    store.upsert_host("host-a", address=None, last_seen_at=datetime(2026, 1, 2, tzinfo=UTC))

    host = store.get_host("host-a")
    assert host is not None
    assert host.address == "10.0.0.1:8000"
    assert host.last_seen_at == datetime(2026, 1, 2, tzinfo=UTC)


def test_capability_tags_round_trip(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    from datetime import UTC, datetime

    store.upsert_host("host-a", address="a", last_seen_at=datetime(2026, 1, 1, tzinfo=UTC))
    store.replace_host_devices(
        "host-a",
        [_pooled("dev-1", "host-a", tags=["ios", "physical"])],
    )

    device = store.get_device("dev-1")
    assert device is not None
    assert device.capability_tags == ["ios", "physical"]


def test_two_hosts_with_colliding_local_device_id_do_not_crash(tmp_path) -> None:
    """Two hosts legitimately reporting the same local device_id string are
    different devices in the pool; syncing both must not violate a uniqueness
    constraint and must keep both rows independently addressable by host."""
    store = CloudStore(tmp_path / "cloud.sqlite3")
    from datetime import UTC, datetime

    store.upsert_host("host-a", address="a", last_seen_at=datetime(2026, 1, 1, tzinfo=UTC))
    store.upsert_host("host-b", address="b", last_seen_at=datetime(2026, 1, 1, tzinfo=UTC))

    store.replace_host_devices("host-a", [_pooled("dev-1", "host-a", status="idle")])
    # Same local device_id "dev-1", different host -- must not raise.
    store.replace_host_devices("host-b", [_pooled("dev-1", "host-b", status="busy")])

    devices = store.list_devices()
    assert len(devices) == 2
    by_host = {d.host_id: d for d in devices}
    assert set(by_host) == {"host-a", "host-b"}
    assert by_host["host-a"].device_id == "dev-1"
    assert by_host["host-b"].device_id == "dev-1"
    assert by_host["host-a"].status == "idle"
    assert by_host["host-b"].status == "busy"

    # Lookup by id alone must still work (finds a device on some host) without crashing.
    found = store.get_device("dev-1")
    assert found is not None
    assert found.host_id in {"host-a", "host-b"}

    # Re-syncing host-a alone must not disturb host-b's colliding-id row.
    store.replace_host_devices("host-a", [_pooled("dev-1", "host-a", status="error")])
    devices = store.list_devices()
    assert len(devices) == 2
    by_host = {d.host_id: d for d in devices}
    assert by_host["host-a"].status == "error"
    assert by_host["host-b"].status == "busy"