model-platform/schedule/tests/test_janitor.py

"""Tests for the node-run janitor (P0-4 S2).

Pre-fix, ``job.node.execute`` outbox retries exhausted (5 tries, capped
30s backoff) marked the *event* ``failed`` but no one reconciled the
*node_run* — the row stayed in ``queued`` / ``running`` forever and
the DAG stopped advancing. The janitor enforces a wall-clock deadline
so any stuck row is eventually flipped to ``timed_out``.

These tests mock the session — no live MySQL. The orchestration
*logic* (which rows count as stuck, what terminal state to write,
which outbox event to enqueue) is what's worth pinning; the SQL itself
is exercised in the integration environment.
"""

from __future__ import annotations

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

import pytest

from common.db.models import ScheduleNodeRuns
from common.eventing import utcnow
from schedule.scheduling.orchestrator import NODE_FINISHED_EVENT, DispatchOrchestrator


def _make_orchestrator() -> DispatchOrchestrator:
    """Construct a ``DispatchOrchestrator`` with a fake session factory."""
    return DispatchOrchestrator(
        session_factory=MagicMock(name="session_factory"),
        node_execute_handler=AsyncMock(),
        execution_concurrency=4,
    )


def _make_node_run(
    *,
    node_run_id: str = "nr1",
    run_id: str = "r1",
    node_id: str = "n1",
    versions_id: str = "v1",
    attempt_no: int = 1,
    node_status: str = "running",
    started_at=None,
    created_at=None,
    finished_at=None,
    state_version: int = 0,
) -> ScheduleNodeRuns:
    """Build an in-memory ``ScheduleNodeRuns`` row."""
    return ScheduleNodeRuns(
        node_run_id=node_run_id,
        run_id=run_id,
        node_id=node_id,
        versions_id=versions_id,
        attempt_no=attempt_no,
        node_status=node_status,
        state_version=state_version,
        created_at=created_at or utcnow(),
        started_at=started_at,
        finished_at=finished_at,
    )


def _make_run(workspace_id: str = "ws1") -> SimpleNamespace:
    return SimpleNamespace(workspace_id=workspace_id)


def _make_node(
    *,
    timeout_seconds: int = 600,
    retry_count: int = 0,
    retry_interval_sec: int = 5,
) -> SimpleNamespace:
    return SimpleNamespace(
        timeout_seconds=timeout_seconds,
        retry_count=retry_count,
        retry_interval_sec=retry_interval_sec,
    )


def _open_session_scope(fake_session: AsyncMock) -> MagicMock:
    """Wrap ``fake_session`` in an async context-manager stand-in."""
    cm = MagicMock()
    cm.__aenter__ = AsyncMock(return_value=fake_session)
    cm.__aexit__ = AsyncMock(return_value=None)
    return cm


async def test_force_terminal_writes_timed_out_and_enqueues_event() -> None:
    """A running row over deadline must flip to ``timed_out`` + emit
    ``NODE_FINISHED_EVENT`` whose idempotency key is the janitor
    variant (``...:timed_out``) so it can't collide with the worker's
    ``...:finished`` key on the same row."""
    orch = _make_orchestrator()
    node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10_000))

    fake_session = AsyncMock(name="session")
    # Re-read under FOR UPDATE returns the same still-running row.
    fake_session.scalar = AsyncMock(return_value=node_run)
    fake_session.add = MagicMock()

    result = await orch._force_terminal_node_run(
        fake_session,
        node_run=node_run,
        run=_make_run(),
        reason="NODE_RUN_TIMEOUT",
    )

    assert result is True
    assert node_run.node_status == "timed_out"
    assert node_run.error_code == "NODE_RUN_TIMEOUT"
    assert node_run.finished_at is not None
    assert node_run.state_version == 1
    assert node_run.duration_ms is not None and node_run.duration_ms >= 0

    fake_session.add.assert_called_once()
    event = fake_session.add.call_args[0][0]
    assert event.event_type == NODE_FINISHED_EVENT
    assert event.aggregate_id == node_run.node_run_id
    assert event.idempotency_key.endswith(":timed_out")
    assert event.payload_json["workspace_id"] == "ws1"
    assert event.payload_json["node_status"] == "timed_out"
    assert event.payload_json["error_code"] == "NODE_RUN_TIMEOUT"
    assert event.payload_json["exit_code"] is None


async def test_force_terminal_skips_when_worker_raced_to_terminal() -> None:
    """If the worker reports a result between SELECT and the lock
    re-read, the janitor must NOT overwrite the worker's terminal
    state — ``_advance_run`` already handled the row correctly."""
    orch = _make_orchestrator()
    candidate = _make_node_run(node_status="running")
    # Re-read under FOR UPDATE shows the worker won the race.
    already_terminal = _make_node_run(node_status="succeeded")

    fake_session = AsyncMock()
    fake_session.scalar = AsyncMock(return_value=already_terminal)
    fake_session.add = MagicMock()

    result = await orch._force_terminal_node_run(
        fake_session,
        node_run=candidate,
        run=_make_run(),
        reason="NODE_RUN_TIMEOUT",
    )

    assert result is False
    fake_session.add.assert_not_called()
    # The row that was re-read under the lock kept its terminal status.
    assert already_terminal.node_status == "succeeded"


async def test_force_terminal_handles_never_started_old_row() -> None:
    """``started_at IS NULL`` + very-old ``created_at`` = row that
    never even got dispatched. Janitor must still force-terminal with
    the "never dispatched" message variant so operators can tell the
    two failure modes apart in logs."""
    orch = _make_orchestrator()
    node_run = _make_node_run(started_at=None, node_status="queued")

    fake_session = AsyncMock()
    fake_session.scalar = AsyncMock(return_value=node_run)
    fake_session.add = MagicMock()

    await orch._force_terminal_node_run(
        fake_session,
        node_run=node_run,
        run=_make_run(),
        reason="NODE_RUN_TIMEOUT",
    )

    assert node_run.node_status == "timed_out"
    assert "从未派发" in (node_run.message or "")
    assert node_run.duration_ms is not None and node_run.duration_ms >= 0


async def test_reap_kills_running_row_past_deadline() -> None:
    """End-to-end of ``_reap_stuck_node_runs``: a row whose
    ``started_at`` is way past ``timeout + grace`` must be flipped,
    and the loop must report ``killed == 1``."""
    orch = _make_orchestrator()
    node = _make_node(timeout_seconds=600)
    node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10_000))
    run = _make_run()

    fake_session = AsyncMock()
    # SELECT returns one candidate.
    fake_session.execute = AsyncMock(
        return_value=MagicMock(all=MagicMock(return_value=[(node_run, node, run)])),
    )
    # Re-read under FOR UPDATE returns the same row.
    fake_session.scalar = AsyncMock(return_value=node_run)
    fake_session.add = MagicMock()

    with patch(
        "schedule.scheduling.orchestrator.session_scope",
        return_value=_open_session_scope(fake_session),
    ):
        killed = await orch._reap_stuck_node_runs()

    assert killed == 1
    assert node_run.node_status == "timed_out"
    fake_session.add.assert_called_once()


async def test_reap_skips_healthy_row() -> None:
    """A row whose deadline is comfortably in the future must not be
    re-read or mutated — the per-row budget check rejects it before
    touching the lock path."""
    orch = _make_orchestrator()
    node = _make_node(timeout_seconds=600)
    # Started 10s ago; deadline is 600 + 120 = 720s from start.
    node_run = _make_node_run(started_at=utcnow() - timedelta(seconds=10))
    run = _make_run()

    fake_session = AsyncMock()
    fake_session.execute = AsyncMock(
        return_value=MagicMock(all=MagicMock(return_value=[(node_run, node, run)])),
    )
    fake_session.scalar = AsyncMock()
    fake_session.add = MagicMock()

    with patch(
        "schedule.scheduling.orchestrator.session_scope",
        return_value=_open_session_scope(fake_session),
    ):
        killed = await orch._reap_stuck_node_runs()

    assert killed == 0
    fake_session.scalar.assert_not_called()
    fake_session.add.assert_not_called()
    assert node_run.node_status == "running"


async def test_reap_processes_multiple_rows_in_one_pass() -> None:
    """A batch of 3 rows — 2 over deadline, 1 healthy — must produce
    exactly 2 kills. The healthy one is short-circuited in Python
    before the lock path."""
    orch = _make_orchestrator()
    node = _make_node(timeout_seconds=600)
    stuck_a = _make_node_run(
        node_run_id="stuck_a", started_at=utcnow() - timedelta(seconds=10_000),
    )
    stuck_b = _make_node_run(
        node_run_id="stuck_b", started_at=utcnow() - timedelta(seconds=9_000),
    )
    healthy = _make_node_run(
        node_run_id="healthy", started_at=utcnow() - timedelta(seconds=10),
    )
    run = _make_run()

    fake_session = AsyncMock()
    fake_session.execute = AsyncMock(
        return_value=MagicMock(
            all=MagicMock(return_value=[
                (stuck_a, node, run),
                (stuck_b, node, run),
                (healthy, node, run),
            ]),
        ),
    )
    # Re-read for the two stuck rows; healthy is skipped before this.
    fake_session.scalar = AsyncMock(side_effect=[stuck_a, stuck_b])
    fake_session.add = MagicMock()

    with patch(
        "schedule.scheduling.orchestrator.session_scope",
        return_value=_open_session_scope(fake_session),
    ):
        killed = await orch._reap_stuck_node_runs()

    assert killed == 2
    assert stuck_a.node_status == "timed_out"
    assert stuck_b.node_status == "timed_out"
    assert healthy.node_status == "running"
    assert fake_session.add.call_count == 2


async def test_janitor_loop_propagates_cancellation() -> None:
    """``_janitor_loop`` must re-raise ``CancelledError`` so ``close()``
    can shut it down cleanly. Patches ``asyncio.sleep`` to raise on
    first call so the test doesn't run forever."""
    orch = _make_orchestrator()
    orch._reap_stuck_node_runs = AsyncMock(return_value=0)

    async def cancel_on_sleep(_seconds: float) -> None:
        raise asyncio.CancelledError

    with patch(
        "schedule.scheduling.orchestrator.asyncio.sleep",
        side_effect=cancel_on_sleep,
    ):
        with pytest.raises(asyncio.CancelledError):
            await orch._janitor_loop()


async def test_janitor_loop_continues_after_reap_exception() -> None:
    """A failure inside one reap pass must not kill the loop — the
    next iteration still runs. The janitor must be self-healing so a
    bad row or transient DB blip can't take the whole service down."""
    orch = _make_orchestrator()
    call_count = 0

    async def _reap_then_succeed() -> int:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            raise RuntimeError("simulated DB blip")
        return 3

    orch._reap_stuck_node_runs = AsyncMock(side_effect=_reap_then_succeed)

    sleep_calls = 0

    async def _count_sleeps(_seconds: float) -> None:
        nonlocal sleep_calls
        sleep_calls += 1
        if sleep_calls >= 2:
            raise asyncio.CancelledError

    with patch(
        "schedule.scheduling.orchestrator.asyncio.sleep",
        side_effect=_count_sleeps,
    ):
        with pytest.raises(asyncio.CancelledError):
            await orch._janitor_loop()

    assert call_count == 2