test_agent_system.py

import json
import os
import tempfile
from pathlib import Path

import pytest

from agent_system import (
    DAGScheduler,
    KeyPool,
    ThrottleController,
    ExecutorAgent,
    Monitor,
    ResourceManager,
    Task,
    TaskStatus,
    AgentRole,
    StateManager,
)


class TestDAGScheduler:
    def test_empty(self):
        dag = DAGScheduler()
        assert dag.get_ready(set()) == []

    def test_single_node(self):
        dag = DAGScheduler()
        dag.add_node("a")
        assert dag.get_ready(set()) == ["a"]
        assert dag.get_ready({"a"}) == []

    def test_dependency_order(self):
        dag = DAGScheduler()
        dag.add_node("a")
        dag.add_node("b", deps=["a"])
        dag.add_node("c", deps=["a"])
        dag.add_node("d", deps=["b", "c"])
        assert dag.get_ready(set()) == ["a"]
        assert dag.get_ready({"a"}) == ["b", "c"]
        assert dag.get_ready({"a", "b", "c"}) == ["d"]
        assert dag.get_ready({"a", "b", "c", "d"}) == []

    def test_to_dict(self):
        dag = DAGScheduler()
        dag.add_node("x", deps=["y"])
        d = dag.to_dict()
        assert "x" in d
        assert d["x"]["deps"] == ["y"]


class TestKeyPool:
    @pytest.mark.asyncio
    async def test_single_key(self):
        pool = KeyPool(["key1"])
        assert await pool.get_key() == "key1"
        assert await pool.get_key() == "key1"

    @pytest.mark.asyncio
    async def test_round_robin(self):
        pool = KeyPool(["k1", "k2"])
        keys = {await pool.get_key() for _ in range(4)}
        assert keys == {"k1", "k2"}

    @pytest.mark.asyncio
    async def test_exhaustion(self):
        pool = KeyPool(["bad"])
        for _ in range(3):
            pool.record_failure("bad")
        assert await pool.get_key() is None
        assert pool.is_exhausted()

    @pytest.mark.asyncio
    async def test_recovery(self):
        pool = KeyPool(["k"])
        pool.record_failure("k")
        pool.record_success("k")
        assert await pool.get_key() == "k"

    def test_to_dict(self):
        pool = KeyPool(["a", "b"])
        d = pool.to_dict()
        assert d["key_count"] == 2
        assert d["available"] == 2


class TestThrottleController:
    def test_acquire_below_limit(self):
        t = ThrottleController(rpm_limit=10)
        wait = t.acquire()
        assert wait == 0.0

    def test_usage_pct(self):
        t = ThrottleController(rpm_limit=100)
        for _ in range(50):
            t.acquire()
        assert 40 <= t.usage_pct() <= 60


class TestExecutorAgent:
    def setup_method(self):
        self.executor = ExecutorAgent()
        self.tmpdir = Path(tempfile.mkdtemp())

    @pytest.mark.asyncio
    async def test_write_file(self):
        patch = {"action": "write", "file": "hello.txt", "content": "world"}
        result = await self.executor.execute(patch, str(self.tmpdir))
        assert result["status"] == "ok"
        assert (self.tmpdir / "hello.txt").read_text() == "world"

    @pytest.mark.asyncio
    async def test_patch_file(self):
        f = self.tmpdir / "test.py"
        f.write_text("old content")
        patch = {"action": "patch", "file": str(f), "old": "old", "new": "new"}
        result = await self.executor.execute(patch, str(self.tmpdir))
        assert result["status"] == "ok"
        assert f.read_text() == "new content"

    @pytest.mark.asyncio
    async def test_patch_not_found(self):
        patch = {"action": "patch", "file": "/nonexistent/file.txt", "old": "a", "new": "b"}
        result = await self.executor.execute(patch)
        assert result["status"] == "error"

    @pytest.mark.asyncio
    async def test_unknown_action(self):
        result = await self.executor.execute({"action": "fly"})
        assert result["status"] == "error"

    @pytest.mark.asyncio
    async def test_append_file(self):
        f = self.tmpdir / "log.txt"
        f.write_text("line1\n")
        patch = {"action": "append", "file": str(f), "content": "line2"}
        await self.executor.execute(patch, str(self.tmpdir))
        assert "line2" in f.read_text()

    @pytest.mark.asyncio
    async def test_delete_file(self):
        f = self.tmpdir / "temp.txt"
        f.write_text("data")
        patch = {"action": "delete", "file": str(f)}
        result = await self.executor.execute(patch, str(self.tmpdir))
        assert result["status"] == "ok"
        assert not f.exists()

    @pytest.mark.asyncio
    async def test_mkdir(self):
        d = self.tmpdir / "newdir" / "sub"
        patch = {"action": "mkdir", "path": str(d)}
        result = await self.executor.execute(patch)
        assert result["status"] == "ok"
        assert d.exists()


class TestMonitor:
    def test_record_and_summary(self):
        m = Monitor()
        m.record("latency", 1.5, "test")
        s = m.summary()
        assert "latency" in s["metrics"]

    def test_alert_rule(self):
        m = Monitor()
        m.add_alert_rule("high_cpu", "cpu", 80, "gt")
        m.record("cpu", 90)
        alerts = m.check_alerts()
        assert len(alerts) == 1
        assert alerts[0]["alert"] == "high_cpu"

    def test_alert_not_triggered(self):
        m = Monitor()
        m.add_alert_rule("high_cpu", "cpu", 80, "gt")
        m.record("cpu", 50)
        assert m.check_alerts() == []


class TestDAGCycleDetection:
    def test_no_cycle(self):
        dag = DAGScheduler()
        dag.add_node("a")
        dag.add_node("b", deps=["a"])
        dag.add_node("c", deps=["b"])
        assert not dag.has_cycle()
        assert dag.validate() == []

    def test_self_cycle(self):
        dag = DAGScheduler()
        dag.add_node("a", deps=["a"])
        assert dag.has_cycle()
        assert "cycle" in dag.validate()[0]

    def test_indirect_cycle(self):
        dag = DAGScheduler()
        dag.add_node("a", deps=["b"])
        dag.add_node("b", deps=["c"])
        dag.add_node("c", deps=["a"])
        assert dag.has_cycle()

    def test_missing_dep(self):
        dag = DAGScheduler()
        dag.add_node("a", deps=["nonexistent"])
        errors = dag.validate()
        assert any("nonexistent" in e for e in errors)

    def test_priority_ordering(self):
        dag = DAGScheduler()
        dag.add_node("low", deps=[], priority=0)
        dag.add_node("high", deps=[], priority=10)
        dag.add_node("mid", deps=[], priority=5)
        ready = dag.get_ready(set())
        assert ready == ["high", "mid", "low"]

    def test_to_dict_includes_priority(self):
        dag = DAGScheduler()
        dag.add_node("x", deps=["y"], priority=5)
        d = dag.to_dict()
        assert d["x"]["priority"] == 5


class TestResourceManager:
    def test_effective_concurrency_default(self):
        rm = ResourceManager(max_concurrency=5)
        assert rm.effective_concurrency() <= 5
        assert rm.effective_concurrency() >= 1

    def test_summary_returns_keys(self):
        rm = ResourceManager()
        s = rm.summary()
        for k in ("cpu_pct", "mem_pct", "effective_concurrency", "skip_review"):
            assert k in s


class TestTask:
    def test_task_creation(self):
        t = Task(id="t1", dag_id="d1", prompt="do something", role=AgentRole.COMMANDER)
        assert t.status == TaskStatus.PENDING
        assert t.id == "t1"

    def test_task_roundtrip(self):
        t = Task(id="t1", dag_id="d1", prompt="hello", role=AgentRole.EXECUTOR,
                 output="done", patch={"action": "write", "file": "x.txt"})
        d = t.to_dict()
        t2 = Task.from_dict(d)
        assert t2.id == t.id
        assert t2.role == t.role
        assert t2.patch == t.patch
        assert t2.output == t.output