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