stack/tests/prisma/test_fleet.py

"""Tests for prisma.fleet.run_pool — the shared multi-host driver (#655)."""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field

from prisma.fleet import MAX_CONSECUTIVE_FAILURES, run_pool
from prisma.llm import LLMCall, LLMMessage, LLMResult


def _call(item: int) -> LLMCall:
    return LLMCall(messages=[LLMMessage(role="user", content=str(item))])


@dataclass
class FakeProvider:
    """Records which host handled each call; optionally fails N times."""

    host: str
    log: list
    lock: threading.Lock
    delay: float = 0.0
    fail_times: int = 0
    calls: int = field(default=0, init=False)

    def complete(self, call: LLMCall) -> LLMResult:
        with self.lock:
            self.calls += 1
            n = self.calls
            self.log.append(self.host)
        if self.delay:
            time.sleep(self.delay)
        if n <= self.fail_times:
            raise RuntimeError(f"{self.host} boom #{n}")
        return LLMResult(text="", tool_calls=[{"name": "x", "input": {}}], usage={})


class TestFanOut:
    def test_uses_both_hosts(self):
        log: list[str] = []
        providers = {
            "A": FakeProvider(host="A", log=log, lock=threading.Lock()),
            "B": FakeProvider(host="B", log=log, lock=threading.Lock()),
        }
        applied: list[int] = []
        errors: list[tuple[int, str]] = []

        run_pool(
            list(range(6)),
            providers,
            per_host=1,
            prepare=_call,
            apply=lambda item, result: applied.append(item),
            on_error=lambda item, exc, stage: errors.append((item, stage)),
        )

        assert errors == []
        assert sorted(applied) == list(range(6))
        assert set(log) == {"A", "B"}


class TestPerHostCap:
    def test_never_exceeds_per_host_in_flight(self):
        per_host = 2
        lock = threading.Lock()
        current = {"A": 0, "B": 0}
        max_seen = {"A": 0, "B": 0}
        log: list[str] = []

        class TrackingProvider:
            def __init__(self, host):
                self.host = host

            def complete(self, call):
                with lock:
                    current[self.host] += 1
                    max_seen[self.host] = max(max_seen[self.host], current[self.host])
                    log.append(self.host)
                time.sleep(0.02)
                with lock:
                    current[self.host] -= 1
                return LLMResult(
                    text="", tool_calls=[{"name": "x", "input": {}}], usage={}
                )

        providers = {"A": TrackingProvider("A"), "B": TrackingProvider("B")}
        applied: list[int] = []

        run_pool(
            list(range(8)),
            providers,
            per_host=per_host,
            prepare=_call,
            apply=lambda item, result: applied.append(item),
            on_error=lambda item, exc, stage: None,
        )

        assert len(applied) == 8
        assert max_seen["A"] <= per_host
        assert max_seen["B"] <= per_host
        assert set(log) == {"A", "B"}


class TestFailingHostRetired:
    def test_dead_host_loses_only_in_flight_items_rest_go_to_survivor(self):
        log: list[str] = []
        lock = threading.Lock()
        # A always fails (and is instant); B always succeeds but is
        # deliberately slower so A's failures resolve first and A gets
        # retired well before B drains the queue.
        providers = {
            "A": FakeProvider(host="A", log=log, lock=lock, fail_times=10**6),
            "B": FakeProvider(host="B", log=log, lock=lock, delay=0.05),
        }
        applied: list[int] = []
        errors: list[tuple[int, str]] = []

        run_pool(
            list(range(5)),
            providers,
            per_host=1,
            prepare=_call,
            apply=lambda item, result: applied.append(item),
            on_error=lambda item, exc, stage: errors.append((item, stage)),
        )

        a_calls = log.count("A")
        b_calls = log.count("B")
        # A is retired after exactly MAX_CONSECUTIVE_FAILURES calls —
        # never dispatched a third time.
        assert a_calls == MAX_CONSECUTIVE_FAILURES
        # Every item not sent to A was resolved by B.
        assert b_calls == 5 - a_calls
        assert len(applied) == b_calls
        assert len(errors) == a_calls
        assert all(stage == "complete" for _, stage in errors)


class TestPrepareSkipsLLM:
    def test_none_from_prepare_skips_dispatch(self):
        providers = {"A": FakeProvider(host="A", log=[], lock=threading.Lock())}
        handled_directly: list[int] = []

        def prepare(item):
            if item % 2 == 0:
                handled_directly.append(item)
                return None
            return _call(item)

        applied: list[int] = []
        run_pool(
            list(range(4)),
            providers,
            per_host=1,
            prepare=prepare,
            apply=lambda item, result: applied.append(item),
            on_error=lambda item, exc, stage: None,
        )

        assert handled_directly == [0, 2]
        assert sorted(applied) == [1, 3]


class TestErrorStages:
    def test_prepare_error_tagged(self):
        providers = {"A": FakeProvider(host="A", log=[], lock=threading.Lock())}
        errors = []

        def prepare(item):
            raise ValueError("bad item")

        run_pool(
            [1],
            providers,
            per_host=1,
            prepare=prepare,
            apply=lambda item, result: None,
            on_error=lambda item, exc, stage: errors.append(stage),
        )
        assert errors == ["prepare"]

    def test_apply_error_tagged(self):
        providers = {"A": FakeProvider(host="A", log=[], lock=threading.Lock())}
        errors = []

        def apply(item, result):
            raise ValueError("bad apply")

        run_pool(
            [1],
            providers,
            per_host=1,
            prepare=_call,
            apply=apply,
            on_error=lambda item, exc, stage: errors.append(stage),
        )
        assert errors == ["apply"]