model-platform/runtime/tests/test_process.py

"""Tests for ``runtime.process`` P0-4 R1 + R2 fixes.

R1: ``_reap_once`` must verify a stuck dead-process record's identity
before deleting it — otherwise a ``start_workspace`` that replaced the
dead record mid-cycle would see its new entry silently erased by the
reaper's stale ``del``.

R2: ``WORKSPACE_LOCKS`` entries must survive ``stop_workspace`` —
popping the lock breaks mutual exclusion for any coroutine still
holding a reference to the old lock object.
"""

from __future__ import annotations

import asyncio
import time
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from runtime import process


def _make_dead_process(pid: int) -> MagicMock:
    """A ``subprocess.Popen``-shaped mock whose ``poll()`` says dead."""
    proc = MagicMock(name=f"proc-{pid}")
    proc.pid = pid
    proc.poll = MagicMock(return_value=0)  # exited
    return proc


def _make_alive_process(pid: int) -> MagicMock:
    """A ``subprocess.Popen``-shaped mock whose ``poll()`` says alive."""
    proc = MagicMock(name=f"proc-{pid}")
    proc.pid = pid
    proc.poll = MagicMock(return_value=None)  # running
    return proc


def _make_record(
    *,
    pid: int,
    started_at: float,
    last_used_at: float | None = None,
    alive: bool = True,
) -> dict:
    """A ``JupyterProcessRecord``-shaped dict (subscriptable, like the TypedDict)."""
    return {
        "process": _make_alive_process(pid) if alive else _make_dead_process(pid),
        "port": 8000 + pid,
        "token": f"token-{pid}",
        "base_url": "http://localhost",
        "started_at": started_at,
        "last_used_at": last_used_at if last_used_at is not None else started_at,
        "meta_path": f"/tmp/{pid}.json",
    }


@pytest.fixture(autouse=True)
def _isolate_registries():
    """Snapshot & restore the global registries so tests don't leak."""
    saved_processes = dict(process.JUPYTER_PROCESSES)
    saved_locks = dict(process.WORKSPACE_LOCKS)
    yield
    process.JUPYTER_PROCESSES.clear()
    process.JUPYTER_PROCESSES.update(saved_processes)
    process.WORKSPACE_LOCKS.clear()
    process.WORKSPACE_LOCKS.update(saved_locks)


# ─── R1: reaper identity check ─────────────────────────────────────────


def test_reap_drops_dead_process_when_identity_matches() -> None:
    """The original (non-buggy) case still works: a dead process with
    nobody replacing it must be cleaned up."""
    now = time.time()
    dead = _make_record(pid=100, started_at=now, alive=False)
    process.JUPYTER_PROCESSES["ws1"] = dead

    with patch.object(process, "_delete_meta") as mock_delete, \
         patch.object(process, "stop_workspace", new=AsyncMock()):
        asyncio.run(process._reap_once())

    assert "ws1" not in process.JUPYTER_PROCESSES
    mock_delete.assert_called_once_with("ws1")


def test_reap_skips_dead_process_when_record_replaced_by_new_process() -> None:
    """P0-4 R1 core bug. ``start_workspace`` detected the dead record
    while the reaper was between snapshot and ``del``, removed it,
    started a new process with a different pid, and wrote the new
    record under the same ``ws_id``. The reaper's stale snapshot now
    references the old (gone) pid; a naive ``del`` would erase the
    new live process. Identity check must catch this and skip."""
    now = time.time()
    replacement = _make_record(pid=200, started_at=now, alive=True)
    # Mutate the dict to simulate ``start_workspace`` having replaced
    # the entry between snapshot time and now.
    process.JUPYTER_PROCESSES["ws1"] = replacement

    with patch.object(process, "_delete_meta") as mock_delete, \
         patch.object(process, "stop_workspace", new=AsyncMock()):
        asyncio.run(process._reap_once())

    # The new live record must survive untouched.
    assert process.JUPYTER_PROCESSES["ws1"] is replacement
    mock_delete.assert_not_called()


def test_reap_keeps_alive_process_alone() -> None:
    """Sanity: a record whose process is still running must not be
    dropped (only marked for idle/max-lifetime stop). With a fresh
    ``started_at`` neither threshold fires, so no victim either."""
    now = time.time()
    live = _make_record(pid=100, started_at=now, alive=True)
    process.JUPYTER_PROCESSES["ws1"] = live

    with patch.object(process, "stop_workspace", new=AsyncMock()) as mock_stop, \
         patch.object(process, "_delete_meta") as mock_delete:
        asyncio.run(process._reap_once())

    assert process.JUPYTER_PROCESSES["ws1"] is live
    mock_delete.assert_not_called()
    mock_stop.assert_not_called()


# ─── R2: workspace locks survive stop ──────────────────────────────────


def test_workspace_lock_is_shared_across_callers() -> None:
    """``get_workspace_lock`` must return the *same* ``asyncio.Lock``
    object every time — that is the whole point of the registry. If
    a caller pops the entry between two ``get_workspace_lock`` calls,
    a fresh lock would be returned and the two callers would no longer
    be serialized."""
    lock_a = process.get_workspace_lock("ws1")
    lock_b = process.get_workspace_lock("ws1")
    assert lock_a is lock_b


def test_concurrent_acquires_serialize_through_same_lock() -> None:
    """Two concurrent callers acquiring the lock for the same
    ``ws_id`` must serialize — proving they share a single lock object,
    not two independent ones. With P0-4 R2 in place (lock never
    popped) this is automatic; before the fix a popped entry would
    let the second caller acquire a fresh, independent lock."""
    lock = process.get_workspace_lock("ws1")
    order: list[str] = []

    async def first() -> None:
        async with lock:
            order.append("first-enter")
            await asyncio.sleep(0)
            order.append("first-exit")

    async def second() -> None:
        # Acquire *after* first releases — if we got a fresh lock here,
        # this would interleave instead of waiting.
        async with lock:
            order.append("second-enter")
            order.append("second-exit")

    async def scenario() -> None:
        await asyncio.gather(first(), second())

    asyncio.run(scenario())
    assert order == [
        "first-enter", "first-exit", "second-enter", "second-exit",
    ], (
        f"Lock leaked: order={order}. The two callers interleaved, "
        "which means they held independent lock objects."
    )


def test_workspace_lock_survives_manual_pop_simulating_stop() -> None:
    """Pre-fix, ``stop_workspace`` popped the lock from the registry
    *after* the async-with exited. The lock object itself survived
    (the dict entry just disappeared), but the next
    ``get_workspace_lock`` call would create a *new* lock and lose
    serialization for any coroutine still holding the old reference.

    This test simulates the pre-fix sequence manually: after one
    ``async with`` block exits, ``get_workspace_lock`` must still
    return the *same* lock if we don't pop it; if we *did* pop it
    (simulating the old buggy behavior), it would return a fresh
    one. The test asserts the fixed invariant: even after the
    stop-shaped sequence, the lock object is the same one."""
    lock_a = process.get_workspace_lock("ws1")

    async def use_and_release() -> None:
        async with lock_a:
            pass

    asyncio.run(use_and_release())

    # Pre-fix path: ``stop_workspace`` would pop the entry here.
    # Post-fix: the helper is gone, so nothing pops the entry. Verify
    # the invariant the fix preserves.
    assert "ws1" in process.WORKSPACE_LOCKS
    lock_b = process.get_workspace_lock("ws1")
    assert lock_a is lock_b


def test_drop_workspace_lock_helper_is_removed() -> None:
    """Belt-and-suspenders: the buggy helper must not come back by
    accident. ``_drop_workspace_lock`` removal is the entire R2 fix;
    if someone re-adds it, this test fails before the leak does."""
    assert not hasattr(process, "_drop_workspace_lock"), (
        "_drop_workspace_lock was re-introduced — P0-4 R2 fix "
        "regressed. Remove it again."
    )


# Late import to avoid an unused-import lint complaint above.