EvoScientist-Multi/tests/test_agent_loader.py

"""Tests for ``cli/_agent_loader``."""

from __future__ import annotations

import asyncio
import threading

import pytest

from EvoScientist.cli._agent_loader import BackgroundAgentLoader, MCPProgressTracker

# ──────────────────────────────────────────────────────────────────────
# MCPProgressTracker
# ──────────────────────────────────────────────────────────────────────


class TestMCPProgressTracker:
    def test_prime_empty_when_no_config(self, monkeypatch):
        import EvoScientist.mcp as mcp_pkg

        monkeypatch.setattr(mcp_pkg, "load_mcp_config", lambda: {})
        t = MCPProgressTracker()
        t.prime()
        assert t.progress == {}

    def test_prime_seeds_pending_entries(self, monkeypatch):
        import EvoScientist.mcp as mcp_pkg

        monkeypatch.setattr(mcp_pkg, "load_mcp_config", lambda: {"a": {}, "b": {}})
        t = MCPProgressTracker()
        t.prime()
        assert t.progress == {"a": ("pending", ""), "b": ("pending", "")}

    def test_prime_swallows_config_errors(self, monkeypatch):
        import EvoScientist.mcp as mcp_pkg

        def _boom():
            raise RuntimeError("config broken")

        monkeypatch.setattr(mcp_pkg, "load_mcp_config", _boom)
        t = MCPProgressTracker()
        t.prime()
        assert t.progress == {}

    def test_record_maps_events(self):
        t = MCPProgressTracker()
        assert t.record("start", "srv", "") == "pending"
        assert t.record("success", "srv", "5") == "ok"
        assert t.record("error", "srv", "timeout") == "error"
        assert t.record("bogus", "srv", "") is None
        assert t.progress == {"srv": ("error", "timeout")}

    def test_start_does_not_overwrite_existing_state(self):
        t = MCPProgressTracker()
        t.record("success", "srv", "3")
        t.record("start", "srv", "")
        assert t.progress["srv"] == ("ok", "3")

    def test_snapshot_is_independent_copy(self):
        t = MCPProgressTracker()
        t.record("success", "srv", "1")
        snap = t.snapshot()
        t.record("error", "srv", "oops")
        assert snap == [("ok", "1")]

    def test_totals(self):
        t = MCPProgressTracker()
        t.record("start", "a", "")
        t.record("success", "b", "1")
        t.record("error", "c", "boom")
        done, total = t.totals()
        assert (done, total) == (2, 3)


# ──────────────────────────────────────────────────────────────────────
# BackgroundAgentLoader
# ──────────────────────────────────────────────────────────────────────


def _make_loader_fn(agent_value="AGENT", fail_with=None, capture=None):
    """Build a sync loader that records ``on_mcp_progress`` + kwargs."""

    def _loader(*, on_mcp_progress=None, **kwargs):
        if capture is not None:
            capture["on_mcp_progress"] = on_mcp_progress
            capture.setdefault("kwargs", []).append(kwargs)
            if on_mcp_progress is not None:
                on_mcp_progress("start", "srv", "")
                on_mcp_progress("success", "srv", "1")
        if fail_with is not None:
            raise fail_with
        return agent_value

    return _loader


class _GatedThreadLoader:
    """Callable loader that blocks until tests explicitly release it."""

    def __init__(self, agent_value="AGENT", progress_events=()):
        self.agent_value = agent_value
        self.progress_events = tuple(progress_events)
        self.started = threading.Event()
        self.release = threading.Event()
        self.finished = threading.Event()

    def __call__(self, *, on_mcp_progress=None):
        self.started.set()
        self.release.wait(timeout=1)
        try:
            if on_mcp_progress is not None:
                for event in self.progress_events:
                    on_mcp_progress(*event)
            return self.agent_value
        finally:
            self.finished.set()


async def _wait_for_event(event, timeout=1):
    return await asyncio.to_thread(event.wait, timeout)


class TestBackgroundAgentLoaderStart:
    async def test_start_creates_task_and_forwards_kwargs(self):
        captured: dict = {}
        loader = BackgroundAgentLoader(_make_loader_fn(capture=captured))

        loader.start(workspace_dir="/ws", checkpointer="CK")
        assert loader.task is not None
        assert loader.is_pending
        await loader.await_ready()

        assert captured["kwargs"][0] == {"workspace_dir": "/ws", "checkpointer": "CK"}

    async def test_start_bumps_load_id(self):
        loader = BackgroundAgentLoader(_make_loader_fn())

        assert loader._load_id == 0
        loader.start()
        assert loader._load_id == 1
        loader.start()
        assert loader._load_id == 2
        await loader.await_ready()

    async def test_start_cancels_in_flight_prior_task(self):
        blocking = _GatedThreadLoader("LATE")

        loader = BackgroundAgentLoader(blocking)
        loader.start()
        first_task = loader.task
        assert first_task is not None
        assert await _wait_for_event(blocking.started)
        # Supersede immediately; asyncio.to_thread wrapper gets cancelled.
        loader._loader_fn = _make_loader_fn("FRESH")
        loader.start()
        agent = await loader.await_ready()
        assert agent == "FRESH"
        blocking.release.set()
        try:
            await first_task
        except asyncio.CancelledError:
            pass
        assert first_task.cancelled() or first_task.done()


class TestBackgroundAgentLoaderCallbacks:
    async def test_progress_hook_sees_events_in_order(self):
        events: list[tuple[str, str, str]] = []
        loader = BackgroundAgentLoader(
            _make_loader_fn(capture={}),
            on_progress=lambda e, s, d: events.append((e, s, d)),
        )

        loader.start()
        await loader.await_ready()

        assert events == [("start", "srv", ""), ("success", "srv", "1")]

    async def test_stale_progress_events_are_dropped(self):
        """A progress event fired after a newer `start` must not reach the hook."""
        slow_loader = _GatedThreadLoader(
            "slow-agent", progress_events=[("success", "from-slow", "1")]
        )
        seen: list[str] = []

        def fast_loader(*, on_mcp_progress=None):
            if on_mcp_progress is not None:
                on_mcp_progress("success", "from-fast", "1")
            return "fast-agent"

        loader = BackgroundAgentLoader(
            slow_loader, on_progress=lambda e, s, d: seen.append(s)
        )

        loader.start()
        assert await _wait_for_event(slow_loader.started)
        # Loader 1 waits so its progress event fires AFTER load 2 starts.
        loader._loader_fn = fast_loader
        loader.start()
        await loader.await_ready()
        slow_loader.release.set()
        assert await _wait_for_event(slow_loader.finished)

        assert "from-fast" in seen
        assert "from-slow" not in seen

    async def test_success_callback_fires_on_completion(self):
        got = []
        loader = BackgroundAgentLoader(
            _make_loader_fn("MY_AGENT"),
            on_success=lambda a: got.append(a),
        )

        loader.start()
        await loader.await_ready()
        await asyncio.sleep(0)  # let done-callback run

        assert got == ["MY_AGENT"]

    async def test_failure_callback_fires_on_error(self):
        err = RuntimeError("load failed")
        got_failures = []
        got_successes = []
        loader = BackgroundAgentLoader(
            _make_loader_fn(fail_with=err),
            on_success=lambda a: got_successes.append(a),
            on_failure=lambda e: got_failures.append(e),
        )

        loader.start()
        with pytest.raises(RuntimeError, match="load failed"):
            await loader.await_ready()
        await asyncio.sleep(0)

        assert got_failures == [err]
        assert got_successes == []


class TestBackgroundAgentLoaderAwaitReady:
    async def test_returns_cached_agent_without_reawaiting(self):
        captured: dict = {}
        loader = BackgroundAgentLoader(_make_loader_fn("A", capture=captured))

        loader.start()
        assert await loader.await_ready() == "A"
        assert await loader.await_ready() == "A"

        assert len(captured["kwargs"]) == 1

    async def test_raises_if_started_not_called(self):
        loader = BackgroundAgentLoader(_make_loader_fn())

        with pytest.raises(RuntimeError, match="before start"):
            await loader.await_ready()

    async def test_reraises_real_error_on_subsequent_awaits(self):
        """After a failure, ``await_ready`` must keep raising the real exception —
        not the "before start()" sentinel — until ``start`` is called again."""

        def _fail(*, on_mcp_progress=None):
            raise RuntimeError("bad MCP config")

        loader = BackgroundAgentLoader(_fail)

        loader.start()
        with pytest.raises(RuntimeError, match="bad MCP config"):
            await loader.await_ready()
        with pytest.raises(RuntimeError, match="bad MCP config"):
            await loader.await_ready()

    async def test_needs_restart_flags_failed_load_for_retry(self):
        calls = {"n": 0}

        def flaky(*, on_mcp_progress=None):
            calls["n"] += 1
            if calls["n"] == 1:
                raise RuntimeError("first attempt failed")
            return "SECOND"

        loader = BackgroundAgentLoader(flaky)

        assert loader.needs_restart  # never started
        loader.start()
        with pytest.raises(RuntimeError):
            await loader.await_ready()
        assert loader.needs_restart  # failed, caller may retry
        loader.start()
        assert await loader.await_ready() == "SECOND"
        assert not loader.needs_restart  # success → no retry


class TestBackgroundAgentLoaderAdopt:
    def test_adopt_seats_external_agent(self):
        loader = BackgroundAgentLoader(_make_loader_fn())
        loader.adopt("EXTERNAL")
        assert loader.agent == "EXTERNAL"
        assert not loader.is_pending

    async def test_adopt_supersedes_in_flight_load(self):
        """A late background completion must not overwrite an adopted agent."""
        slow_loader = _GatedThreadLoader("FROM_BACKGROUND")

        loader = BackgroundAgentLoader(slow_loader)

        loader.start()
        assert await _wait_for_event(slow_loader.started)
        loader.adopt("FROM_MODEL")
        slow_loader.release.set()
        assert await _wait_for_event(slow_loader.finished)
        await asyncio.sleep(0)
        assert loader.agent == "FROM_MODEL"


class TestBackgroundAgentLoaderIsPending:
    def test_false_before_start(self):
        loader = BackgroundAgentLoader(_make_loader_fn())
        assert not loader.is_pending

    async def test_false_after_completion(self):
        loader = BackgroundAgentLoader(_make_loader_fn())

        loader.start()
        await loader.await_ready()

        assert not loader.is_pending

    async def test_true_between_start_and_completion(self):
        wait_loader = _GatedThreadLoader("ok")

        loader = BackgroundAgentLoader(wait_loader)

        loader.start()
        assert await _wait_for_event(wait_loader.started)
        assert loader.is_pending
        wait_loader.release.set()
        await loader.await_ready()
        assert not loader.is_pending