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

from __future__ import annotations

import threading
from datetime import UTC, datetime

from host_agent.mcp_lock import McpBusyTracker


def _tracker_with_now() -> tuple[McpBusyTracker, list[datetime]]:
    times: list[datetime] = []

    def now() -> datetime:
        return times[-1] if times else datetime(2026, 1, 1, tzinfo=UTC)

    tracker = McpBusyTracker(ttl_seconds=60.0, now=now)
    return tracker, times


def test_acquire_succeeds_on_empty() -> None:
    tracker, _ = _tracker_with_now()
    assert tracker.acquire("phone-1", "sess-a") is True
    assert "phone-1" in tracker.busy_device_ids()


def test_acquire_fails_when_held_by_other_session() -> None:
    tracker, _ = _tracker_with_now()
    assert tracker.acquire("phone-1", "sess-a") is True
    assert tracker.acquire("phone-1", "sess-b") is False


def test_acquire_is_idempotent_for_same_session() -> None:
    tracker, _ = _tracker_with_now()
    assert tracker.acquire("phone-1", "sess-a") is True
    # Same session re-acquiring is allowed (acts as renew).
    assert tracker.acquire("phone-1", "sess-a") is True


def test_renew_refreshes_last_seen() -> None:
    tracker, times = _tracker_with_now()
    times.append(datetime(2026, 1, 1, 12, 0, tzinfo=UTC))
    tracker.acquire("phone-1", "sess-a")
    initial = tracker.snapshot()[0]
    times.append(datetime(2026, 1, 1, 12, 0, 30, tzinfo=UTC))
    assert tracker.renew("phone-1", "sess-a") is True
    refreshed = tracker.snapshot()[0]
    assert refreshed.last_seen_at > initial.last_seen_at


def test_renew_fails_when_held_by_other() -> None:
    tracker, _ = _tracker_with_now()
    tracker.acquire("phone-1", "sess-a")
    assert tracker.renew("phone-1", "sess-b") is False


def test_release_returns_freed_device_ids() -> None:
    tracker, _ = _tracker_with_now()
    tracker.acquire("phone-1", "sess-a")
    tracker.acquire("phone-2", "sess-a")
    freed = tracker.release("sess-a")
    assert sorted(freed) == ["phone-1", "phone-2"]
    assert tracker.busy_device_ids() == []


def test_release_only_frees_caller_session() -> None:
    tracker, _ = _tracker_with_now()
    tracker.acquire("phone-1", "sess-a")
    tracker.acquire("phone-1", "sess-b")  # fails
    freed = tracker.release("sess-b")
    assert freed == []
    assert "phone-1" in tracker.busy_device_ids()


def test_ttl_sweeps_expired_leases() -> None:
    tracker, times = _tracker_with_now()
    times.append(datetime(2026, 1, 1, 12, 0, tzinfo=UTC))
    tracker.acquire("phone-1", "sess-a")
    # Advance past TTL without renew.
    times.append(datetime(2026, 1, 1, 12, 1, 1, tzinfo=UTC))  # 61s later
    assert tracker.busy_device_ids() == []


def test_renew_after_ttl_tolerates_same_session() -> None:
    """Scene 10: lease expired but session_id matches -> re-acquire."""
    tracker, times = _tracker_with_now()
    times.append(datetime(2026, 1, 1, 12, 0, tzinfo=UTC))
    tracker.acquire("phone-1", "sess-a")
    times.append(datetime(2026, 1, 1, 12, 1, 1, tzinfo=UTC))  # expired
    # renew from the same session should succeed (re-acquire).
    assert tracker.renew("phone-1", "sess-a") is True
    assert "phone-1" in tracker.busy_device_ids()


def test_snapshot_matches_busy_device_ids() -> None:
    tracker, _ = _tracker_with_now()
    tracker.acquire("phone-1", "sess-a")
    tracker.acquire("phone-2", "sess-a")
    snap = tracker.snapshot()
    assert {lease.device_id for lease in snap} == set(tracker.busy_device_ids())


def test_wait_until_usable_succeeds_when_free() -> None:
    tracker, _ = _tracker_with_now()
    ok = tracker.wait_until_usable("phone-1", "sess-a", timeout=1.0, poll_interval=0.01)
    assert ok is True
    assert "phone-1" in tracker.busy_device_ids()


def test_wait_until_usable_returns_false_on_timeout() -> None:
    tracker, _ = _tracker_with_now()
    tracker.acquire("phone-1", "sess-a")
    ok = tracker.wait_until_usable("phone-1", "sess-b", timeout=0.1, poll_interval=0.02)
    assert ok is False


def test_wait_until_usable_blocks_then_succeeds_when_released() -> None:
    tracker, _ = _tracker_with_now()
    tracker.acquire("phone-1", "sess-a")

    def releaser() -> None:
        import time

        time.sleep(0.05)
        tracker.release("sess-a")

    t = threading.Thread(target=releaser)
    t.start()
    try:
        ok = tracker.wait_until_usable(
            "phone-1", "sess-b", timeout=2.0, poll_interval=0.02
        )
        assert ok is True
    finally:
        t.join()


def test_wait_until_usable_blocks_then_fails_when_cloud_remains_busy() -> None:
    tracker, _ = _tracker_with_now()
    ok = tracker.wait_until_usable(
        "phone-1",
        "sess-a",
        timeout=0.1,
        poll_interval=0.02,
        cloud_busy_check=lambda: True,
    )
    assert ok is False
    assert tracker.busy_device_ids() == []


def test_default_ttl_is_20_seconds() -> None:
    """The McpBusyTracker default TTL is 20s: short enough to recover from a
    dead MCP session within a heartbeat interval without an explicit release
    callback (mcp SDK 1.28.1 has no per-session shutdown hook), but long
    enough that an actively-busy session does not lose its lease during
    normal operator pauses."""
    tracker = McpBusyTracker()
    assert tracker._ttl == 20.0


def test_default_ttl_recovers_dead_session_within_one_window() -> None:
    """With the 20s default, a session that never renews its lease is
    reaped within one TTL window on the next read. This is the
    concrete fallback behavior for I1 (no FastMCP session-end hook)."""
    times: list[datetime] = []

    def now() -> datetime:
        return times[-1] if times else datetime(2026, 1, 1, tzinfo=UTC)

    tracker = McpBusyTracker(now=now)  # default 20s TTL
    times.append(datetime(2026, 1, 1, 12, 0, tzinfo=UTC))
    assert tracker.acquire("phone-1", "dead-session") is True
    # No renew: advance 21s. Lease should be swept on next read.
    times.append(datetime(2026, 1, 1, 12, 0, 21, tzinfo=UTC))
    assert tracker.busy_device_ids() == []
    # New session can now acquire cleanly (no stale-busy contamination).
    assert tracker.acquire("phone-1", "new-session") is True