app/backend/tests/flow/test_supervision.py

"""Supervision: a background task that dies gets restarted, until it is hopeless."""

import asyncio

import pytest

from fluksio.flow import supervision
from fluksio.flow.events import EventBus
from fluksio.flow.supervision import FAILURE_BUDGET, Supervisor


@pytest.fixture(autouse=True)
def no_backoff(monkeypatch: pytest.MonkeyPatch):
    """Real backoff would make these tests a minute long."""
    monkeypatch.setattr(supervision, "BACKOFF", (0.0,))


async def _settle() -> None:
    """Let the supervisor work through its restarts."""
    for _ in range(50):
        await asyncio.sleep(0)


def test_a_crashing_loop_is_restarted():
    attempts = 0

    async def loop() -> None:
        nonlocal attempts
        attempts += 1
        if attempts < 3:
            raise ConnectionError("broker went away")

    async def scenario() -> None:
        supervisor = Supervisor()
        supervisor.spawn("mqtt", "heating", loop)
        await _settle()

        assert attempts == 3
        assert supervisor.quarantined == set()
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_a_loop_that_returns_is_left_alone():
    """Returning is how a loop says it is finished, not that it failed."""
    attempts = 0

    async def loop() -> None:
        nonlocal attempts
        attempts += 1

    async def scenario() -> None:
        supervisor = Supervisor()
        supervisor.spawn("cron", "heating", loop)
        await _settle()

        assert attempts == 1
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_a_flow_that_keeps_crashing_is_quarantined():
    attempts = 0
    events: list[dict] = []
    bus = EventBus()
    bus.publish = events.append  # type: ignore[method-assign]

    async def loop() -> None:
        nonlocal attempts
        attempts += 1
        raise RuntimeError("nope")

    async def scenario() -> None:
        supervisor = Supervisor(bus)
        supervisor.spawn("mqtt", "heating", loop)
        await _settle()

        # Restarting stops once the budget is gone, rather than spinning forever.
        assert attempts == FAILURE_BUDGET
        assert supervisor.quarantined == {"heating"}
        assert [e["type"] for e in events].count("task_crashed") == FAILURE_BUDGET
        assert [e["type"] for e in events][-1] == "flow_quarantined"
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_the_budget_is_per_flow():
    ran = 0

    async def failing() -> None:
        raise RuntimeError("nope")

    async def healthy() -> None:
        nonlocal ran
        ran += 1
        await asyncio.sleep(3600)

    async def scenario() -> None:
        supervisor = Supervisor()
        supervisor.spawn("a", "broken", failing)
        supervisor.spawn("b", "fine", healthy)
        await _settle()

        assert supervisor.quarantined == {"broken"}
        assert ran == 1
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_cancel_all_is_idempotent():
    async def loop() -> None:
        await asyncio.sleep(3600)

    async def scenario() -> None:
        supervisor = Supervisor()
        supervisor.spawn("mqtt", "heating", loop)
        await _settle()

        await supervisor.cancel_all()
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_cancelling_one_flow_leaves_another_flows_tasks_running():
    """Rebuilding one flow must not take every other flow's loops with it."""

    async def loop() -> None:
        await asyncio.sleep(3600)

    async def scenario() -> None:
        supervisor = Supervisor()
        supervisor.spawn("gone.sub", "gone", loop)
        supervisor.spawn("fine.sub", "fine", loop)
        await _settle()

        survivor = supervisor._tasks["fine.sub"]
        await supervisor.cancel_flow("gone")

        assert list(supervisor._tasks) == ["fine.sub"]
        assert not survivor.done()
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_cancelling_one_flow_lifts_only_its_quarantine():
    async def failing() -> None:
        raise RuntimeError("nope")

    async def scenario() -> None:
        supervisor = Supervisor()
        supervisor.spawn("a", "broken", failing)
        supervisor.spawn("b", "also-broken", failing)
        await _settle()

        assert supervisor.quarantined == {"broken", "also-broken"}

        await supervisor.cancel_flow("broken")

        assert supervisor.quarantined == {"also-broken"}
        # And its budget, or the flow would be quarantined again on the first
        # crash after the rebuild that was meant to fix it.
        assert "broken" not in supervisor._failures
        assert "also-broken" in supervisor._failures
        await supervisor.cancel_all()

    asyncio.run(scenario())