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, QUARANTINE_BACKOFF, 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,))
    # The quarantine keeps its real length, so it holds until a test cancels it
    # however slow the run is; 50ms here ran out mid-assertion on a loaded run.
    # The one test that wants it to expire shortens it itself.


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_stood_down():
    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
        quarantine = [e for e in events if e["type"] == "flow_quarantined"][-1]
        assert quarantine["retry_in_s"] == QUARANTINE_BACKOFF[0]
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_one_flows_failure_does_not_stand_another_down():
    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_two_tasks_on_one_flow_do_not_spend_each_others_budget():
    """The house's power flow: an MQTT subscriber and its keepalive publisher.

    Both die with the same broker. Sharing one budget meant the flow was stood
    down after half the restarts either task was allowed on its own.
    """
    tries = {"sub": 0, "pub": 0}

    def crasher(which: str):
        async def loop() -> None:
            tries[which] += 1
            raise ConnectionError("broker went away")

        return loop

    async def scenario() -> None:
        supervisor = Supervisor()
        supervisor.spawn("power.cerbo:mqtt", "power", crasher("sub"))
        supervisor.spawn("power.keepalive:mqtt-out", "power", crasher("pub"))
        await _settle()

        assert tries == {"sub": FAILURE_BUDGET, "pub": FAILURE_BUDGET}
        await supervisor.cancel_all()

    asyncio.run(scenario())


def test_a_quarantine_expires_and_the_task_tries_again(
    monkeypatch: pytest.MonkeyPatch,
):
    """The whole point: a broker that comes back is picked up without a human."""
    monkeypatch.setattr(supervision, "QUARANTINE_BACKOFF", (0.05,))
    attempts = 0

    async def loop() -> None:
        nonlocal attempts
        attempts += 1
        if attempts <= FAILURE_BUDGET:
            raise ConnectionError("broker went away")
        await asyncio.sleep(3600)

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

        assert attempts == FAILURE_BUDGET
        assert supervisor.quarantined == {"power"}

        # Sit out the rest, shortened to 50ms above.
        await asyncio.sleep(0.1)
        await _settle()

        assert attempts == FAILURE_BUDGET + 1
        assert supervisor.quarantined == set()
        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 "a" not in supervisor._failures
        assert "b" in supervisor._failures
        await supervisor.cancel_all()

    asyncio.run(scenario())