roboco/tests/unit/runtime/test_orchestrator_shutdown_drain.py

"""Drain ``_bg_tasks`` on shutdown so fire-and-forget writes (respawn_tracker
upserts, audit-log writes, intake first-message delivery) are not abandoned.

Invariant: ``Orchestrator.stop()`` drains ``_bg_tasks`` with a bounded timeout —
short DB writes finish before the process exits (data preserved), a stuck task
is cancelled once the deadline passes (can't hang shutdown). The ``stop_agent``
loop is wrapped so one agent's stop error can't skip the drain.
"""

from __future__ import annotations

import asyncio
from typing import Any
from unittest.mock import patch

import pytest
from roboco.models.runtime import AgentInstance
from roboco.runtime.orchestrator import (
    _SHUTDOWN_DRAIN_TIMEOUT_SECONDS,
    AgentOrchestrator,
)

# Floor encoding the logical-regression guard: a drain deadline below this
# would risk dropping a legitimate short DB write (an upsert that needs a
# second under load) before it commits — the exact data loss this fix targets.
# Named (not magic) for ruff PLR2004.
_MIN_DRAIN_TIMEOUT = 3.0


def _make_orchestrator() -> AgentOrchestrator:
    """AgentOrchestrator with constructor I/O skipped; stop() deps ready.

    ``stop()`` cancels the named loop tasks (all None here → no-op) and the
    agents in ``_instances`` (empty here), then must drain ``_bg_tasks``.
    """
    with patch.object(AgentOrchestrator, "__init__", return_value=None):
        orch = AgentOrchestrator.__new__(AgentOrchestrator)
    orch._instances = {}
    orch._bg_tasks = set()
    # Every named background loop ``stop()`` cancels — None makes each a no-op
    # so the test exercises ONLY the _bg_tasks drain.
    for attr in (
        "_health_task",
        "_dispatcher_task",
        "_sweeper_task",
        "_rate_limit_probe_task",
        "_strategy_engine_task",
        "_external_pr_poll_task",
        "_self_heal_task",
        "_ci_watch_task",
        "_dep_update_task",
        "_release_manager_task",
    ):
        setattr(orch, attr, None)
    return orch


def test_shutdown_drain_timeout_is_named_module_constant() -> None:
    assert isinstance(_SHUTDOWN_DRAIN_TIMEOUT_SECONDS, int | float)
    assert _SHUTDOWN_DRAIN_TIMEOUT_SECONDS > 0


def test_shutdown_drain_timeout_is_generous() -> None:
    """A short DB upsert under load can legitimately take a moment; the drain
    deadline must not drop it. This guards the logical regression: a too-short
    drain would silently lose the very writes it exists to preserve."""
    assert _SHUTDOWN_DRAIN_TIMEOUT_SECONDS >= _MIN_DRAIN_TIMEOUT


@pytest.mark.asyncio
async def test_stop_drains_completing_bg_task_before_returning() -> None:
    """A bg task that finishes quickly MUST complete (its side effect observed)
    before ``stop()`` returns. Without the drain, ``stop()`` returns immediately
    and the task is abandoned mid-flight — the data-loss tail."""
    orch = _make_orchestrator()
    ran: list[bool] = []

    async def _completes() -> None:
        await asyncio.sleep(0.01)
        ran.append(True)

    orch._bg_tasks.add(asyncio.create_task(_completes()))

    await asyncio.wait_for(orch.stop(), timeout=5.0)

    assert ran == [True], "completing bg task was abandoned by stop()"


@pytest.mark.asyncio
async def test_stop_does_not_hang_on_stuck_bg_task(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A bg task that never completes MUST NOT hang shutdown past the drain
    deadline — it is cancelled once the drain times out. Without the drain,
    a stuck bg task would let ``stop()`` (and thus the process) hang forever.

    Deterministic: the drain deadline is patched tiny so a bounded fail-close is
    asserted in well under a second, never relying on the real 5s default."""
    monkeypatch.setattr(
        "roboco.runtime.orchestrator._SHUTDOWN_DRAIN_TIMEOUT_SECONDS", 0.05
    )
    orch = _make_orchestrator()

    async def _hangs() -> None:
        await asyncio.Future()  # never resolves

    stuck = asyncio.create_task(_hangs())
    orch._bg_tasks.add(stuck)

    await asyncio.wait_for(orch.stop(), timeout=2.0)

    assert stuck.cancelled(), "stuck bg task was not cancelled by the drain"


@pytest.mark.asyncio
async def test_stop_failing_agent_does_not_skip_drain() -> None:
    """If one agent's ``stop_agent`` raises, the drain must still run —
    otherwise a single bad agent would re-introduce the data-loss tail for every
    in-flight bg write. The completing bg task should still finish."""
    orch = _make_orchestrator()
    orch._instances["bad-agent"] = AgentInstance(agent_id="bad-agent")

    async def _raises(_aid: str, **_kw: Any) -> None:
        raise RuntimeError("boom")

    patch.object(orch, "stop_agent", _raises).start()

    ran: list[bool] = []

    async def _completes() -> None:
        await asyncio.sleep(0.01)
        ran.append(True)

    orch._bg_tasks.add(asyncio.create_task(_completes()))

    await asyncio.wait_for(orch.stop(), timeout=5.0)

    assert ran == [True], "failing stop_agent skipped the drain (data lost)"


@pytest.mark.asyncio
async def test_stop_is_idempotent_double_call_is_noop() -> None:
    """stop() is idempotent: the lifespan path and bootstrap's finally block both
    call it, so the second call must be a clean no-op — not a re-drain or re-stop
    of already-stopped agents — guarded by ``_stopped``."""
    orch = _make_orchestrator()
    real_drain = orch._drain_bg_tasks
    drain_calls = 0

    async def counting_drain() -> None:
        nonlocal drain_calls
        drain_calls += 1
        await real_drain()

    orch._drain_bg_tasks = counting_drain

    await orch.stop()
    assert drain_calls == 1, "first stop() drained the bg tasks"
    assert orch._stopped is True

    await orch.stop()  # safety-net double-call (lifespan already stopped it)
    assert drain_calls == 1, "second stop() must not re-drain (idempotent no-op)"
    assert orch._stopped is True