agentic-mobile-control/tests/test_device_pool.py

"""Unit tests for cloud.pool.DevicePool (task 3.4)."""

from __future__ import annotations

from datetime import UTC, datetime, timedelta

from cloud.config import CloudConfig
from cloud.pool import DevicePool
from cloud.store import CloudStore
from core.models import Device


def _config(**overrides) -> CloudConfig:
    base = {
        "sync_interval_seconds": 30,
        "stale_after_seconds": 60,
        "max_queue_depth": 100,
        "default_assignment_strategy": "fifo_match",
        "api_version_prefix": "/v1",
        "db_path": "cloud/cloud.sqlite3",
    }
    base.update(overrides)
    return CloudConfig(**base)


def _device(
    device_id: str,
    *,
    status: str = "idle",
    driver_type: str = "wda",
    capability_tags: list[str] | None = None,
) -> Device:
    return Device(
        id=device_id,
        status=status,  # type: ignore[arg-type]
        driver_type=driver_type,
        capability_tags=list(capability_tags or []),
    )


def test_new_host_sync_creates_registration_and_devices(tmp_path) -> None:
    pool = DevicePool(
        CloudStore(tmp_path / "cloud.sqlite3"),
        _config(),
    )

    pool.sync_host_devices(
        "host-a",
        [_device("dev-1"), _device("dev-2", status="busy")],
        address="10.0.0.1:8000",
    )

    hosts = pool.list_hosts()
    assert [h.host_id for h in hosts] == ["host-a"]
    assert hosts[0].address == "10.0.0.1:8000"

    devices = pool.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"].status == "idle"
    assert by_id["dev-2"].status == "busy"
    assert all(d.host_id == "host-a" for d in devices)


def test_resync_updates_last_seen_and_replaces_devices(tmp_path) -> None:
    pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())

    pool.sync_host_devices("host-a", [_device("dev-1"), _device("dev-2")])
    first_hosts = pool.list_hosts()
    first_seen = first_hosts[0].last_seen_at

    # Force time forward by directly mutating the stored timestamp.
    pool.store.upsert_host(
        "host-a",
        address=None,
        last_seen_at=datetime.now(UTC) - timedelta(seconds=10),
    )

    pool.sync_host_devices("host-a", [_device("dev-3")])
    hosts = pool.list_hosts()
    devices = pool.list_devices()

    assert [h.host_id for h in hosts] == ["host-a"]
    assert {d.device_id for d in devices} == {"dev-3"}
    assert hosts[0].last_seen_at > first_seen


def test_stale_host_devices_reported_unreachable(tmp_path) -> None:
    pool = DevicePool(
        CloudStore(tmp_path / "cloud.sqlite3"),
        _config(stale_after_seconds=60),
    )
    pool.sync_host_devices("host-a", [_device("dev-1", status="idle")])

    # Push the host's last_seen_at beyond the staleness threshold.
    pool.store.upsert_host(
        "host-a",
        address=None,
        last_seen_at=datetime.now(UTC) - timedelta(seconds=120),
    )

    devices = pool.list_devices()
    assert len(devices) == 1
    assert devices[0].status == "unreachable"

    fetched = pool.get_device("dev-1")
    assert fetched is not None
    assert fetched.status == "unreachable"


def test_resync_after_stale_clears_unreachable(tmp_path) -> None:
    pool = DevicePool(
        CloudStore(tmp_path / "cloud.sqlite3"),
        _config(stale_after_seconds=60),
    )
    pool.sync_host_devices("host-a", [_device("dev-1", status="idle")])
    pool.store.upsert_host(
        "host-a",
        address=None,
        last_seen_at=datetime.now(UTC) - timedelta(seconds=120),
    )

    # Stale right now.
    assert pool.list_devices()[0].status == "unreachable"

    # Host resyncs with a fresh snapshot.
    pool.sync_host_devices("host-a", [_device("dev-1", status="idle")])
    devices = pool.list_devices()
    assert devices[0].status == "idle"


def test_unknown_device_returns_none(tmp_path) -> None:
    pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
    assert pool.get_device("does-not-exist") is None


def test_empty_pool_returns_empty_list(tmp_path) -> None:
    pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
    assert pool.list_devices() == []
    assert pool.list_hosts() == []


def test_two_hosts_aggregate_into_one_listing(tmp_path) -> None:
    pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
    pool.sync_host_devices("host-a", [_device("a-dev-1")], address="a")
    pool.sync_host_devices("host-b", [_device("b-dev-1"), _device("b-dev-2")], address="b")

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


def test_two_hosts_syncing_same_local_device_id_do_not_crash(tmp_path) -> None:
    """Two hosts legitimately reporting the same local device_id are distinct
    devices in the pool; syncing both must succeed and yield two entries."""
    pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
    pool.sync_host_devices("host-a", [_device("dev-1", status="idle")], address="a")
    pool.sync_host_devices("host-b", [_device("dev-1", status="busy")], address="b")

    devices = pool.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"


def test_capability_tags_flow_from_device_to_pooled_device(tmp_path) -> None:
    pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
    pool.sync_host_devices(
        "host-a",
        [_device("dev-1", capability_tags=["ios", "physical"])],
        address="a",
    )

    devices = pool.list_devices()
    assert len(devices) == 1
    assert devices[0].capability_tags == ["ios", "physical"]

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