agentic-mobile-control/apps/device-host-agent/tests/test_heartbeat.py

from __future__ import annotations

import asyncio
from datetime import UTC, datetime

from cloud.internal_api.models import HeartbeatResponse
from cloud.internal_api.models import HostGovernancePolicyModel
from device.manager import DeviceManager
from host_agent.config import HostAgentConfig
from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot
from host_agent.mcp_lock import McpBusyTracker
from host_agent.policy_cache import HostPolicyCacheStore
from host_agent.status import AgentStatusTracker


class ConnectableDriver:
    def connect(self) -> None:
        return None


def _config() -> HostAgentConfig:
    return HostAgentConfig(
        control_plane_url="https://control.example",
        host_id="host-a",
        token="secret",
        heartbeat_interval_seconds=0.01,
    )


def test_build_device_snapshot_copies_complete_non_secret_state() -> None:
    manager = DeviceManager()
    manager.register_device(
        "device-b",
        lambda: object(),  # type: ignore[arg-type,return-value]
        driver_type="appium",
        connection_info={"password": "must-not-leave-host"},
        status="busy",
        capability_tags=["android", "physical"],
    )
    manager.register_device(
        "device-a",
        lambda: object(),  # type: ignore[arg-type,return-value]
        status="offline",
    )

    snapshot = build_device_snapshot(manager)

    assert [device.device_id for device in snapshot] == ["device-a", "device-b"]
    assert snapshot[1].driver_type == "appium"
    assert snapshot[1].status == "idle"
    assert snapshot[1].capability_tags == ["android", "physical"]
    assert "must-not-leave-host" not in repr(snapshot)


def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> None:
    manager = DeviceManager()
    manager.register_device(
        "device-a",
        lambda: ConnectableDriver(),  # type: ignore[arg-type,return-value]
    )
    calls: list[list[str]] = []

    class FakeClient:
        async def heartbeat(
            self, devices, *, address=None, policy_revision=0, **kwargs
        ):
            calls.append([device.device_id for device in devices])
            return HeartbeatResponse(
                host_id="host-a",
                accepted_devices=len(devices),
                received_at=datetime.now(UTC),
            )

    async def scenario() -> None:
        stop = asyncio.Event()
        synchronizer = HeartbeatSynchronizer(
            manager,
            FakeClient(),  # type: ignore[arg-type]
            _config(),
        )
        task = asyncio.create_task(synchronizer.run(stop))
        while len(calls) < 3:
            await asyncio.sleep(0.005)
        stop.set()
        await task

    asyncio.run(scenario())
    assert calls == [["device-a"], ["device-a"], ["device-a"]]
    assert manager.status("device-a") == "busy"


def test_sync_once_notifies_status_tracker_and_on_sync_with_device_count() -> None:
    manager = DeviceManager()
    manager.register_device(
        "device-a",
        lambda: ConnectableDriver(),  # type: ignore[arg-type,return-value]
    )
    manager.register_device(
        "device-b",
        lambda: ConnectableDriver(),  # type: ignore[arg-type,return-value]
    )

    class FakeClient:
        async def heartbeat(
            self, devices, *, address=None, policy_revision=0, **kwargs
        ):
            return HeartbeatResponse(
                host_id="host-a",
                accepted_devices=len(devices),
                received_at=datetime.now(UTC),
            )

    async def scenario() -> None:
        tracker = AgentStatusTracker()
        on_sync_calls: list[int] = []
        synchronizer = HeartbeatSynchronizer(
            manager,
            FakeClient(),  # type: ignore[arg-type]
            _config(),
            status_tracker=tracker,
            on_sync=on_sync_calls.append,
        )

        await synchronizer.sync_once()

        assert on_sync_calls == [2]
        last_heartbeat = tracker.snapshot()["last_heartbeat"]
        assert last_heartbeat is not None
        assert last_heartbeat["ok"] is True
        assert last_heartbeat["device_count"] == 2

    asyncio.run(scenario())


def test_heartbeat_caches_safe_host_policy_and_reuses_its_revision(tmp_path) -> None:
    manager = DeviceManager()
    cache = HostPolicyCacheStore(tmp_path / "host_policy.json")
    revisions: list[int] = []

    class UpdatingClient:
        async def heartbeat(
            self, devices, *, address=None, policy_revision=0, **kwargs
        ):
            revisions.append(policy_revision)
            return HeartbeatResponse(
                host_id="host-a",
                accepted_devices=len(devices),
                received_at=datetime.now(UTC),
                policy_revision=4,
                policy=HostGovernancePolicyModel(
                    revision=4,
                    self_submission_enabled=False,
                    max_active_tasks=2,
                    daily_token_budget=900,
                ),
            )

    async def scenario() -> None:
        tracker = AgentStatusTracker()
        synchronizer = HeartbeatSynchronizer(
            manager,
            UpdatingClient(),  # type: ignore[arg-type]
            _config(),
            policy_cache=cache,
            status_tracker=tracker,
        )
        await synchronizer.sync_once()
        assert tracker.snapshot()["host_policy"] == {
            "revision": 4,
            "self_submission_enabled": False,
            "max_active_tasks": 2,
            "daily_token_budget": 900,
        }

        restarted = HeartbeatSynchronizer(
            manager,
            UpdatingClient(),  # type: ignore[arg-type]
            _config(),
            policy_cache=cache,
        )
        assert restarted.policy_revision == 4

    asyncio.run(scenario())
    assert revisions == [0]
    assert '"token":' not in (tmp_path / "host_policy.json").read_text(encoding="utf-8")


def test_sync_once_passes_mcp_busy_device_ids_to_client() -> None:
    """When mcp_busy_tracker has a lease, sync_once relays the device_ids."""
    manager = DeviceManager()
    tracker = McpBusyTracker()
    assert tracker.acquire("phone-1", "sess-a")
    last_kwargs: dict[str, object] = {}

    class FakeClient:
        async def heartbeat(
            self, devices, *, address=None, policy_revision=0, **kwargs
        ):
            last_kwargs.update(kwargs)
            return HeartbeatResponse(
                host_id="host-a",
                accepted_devices=len(devices),
                received_at=datetime.now(UTC),
            )

    async def scenario() -> None:
        sync = HeartbeatSynchronizer(
            manager,
            FakeClient(),  # type: ignore[arg-type]
            _config(),
            mcp_busy_tracker=tracker,
        )
        await sync.sync_once()

    asyncio.run(scenario())
    assert last_kwargs.get("mcp_busy_device_ids") == ["phone-1"]


def test_sync_once_passes_empty_when_tracker_is_none() -> None:
    """Default: no tracker → no busy device ids forwarded."""
    manager = DeviceManager()
    last_kwargs: dict[str, object] = {}

    class FakeClient:
        async def heartbeat(
            self, devices, *, address=None, policy_revision=0, **kwargs
        ):
            last_kwargs.update(kwargs)
            return HeartbeatResponse(
                host_id="host-a",
                accepted_devices=len(devices),
                received_at=datetime.now(UTC),
            )

    async def scenario() -> None:
        sync = HeartbeatSynchronizer(
            manager,
            FakeClient(),  # type: ignore[arg-type]
            _config(),
        )
        await sync.sync_once()

    asyncio.run(scenario())
    assert not last_kwargs.get("mcp_busy_device_ids")