Files
2026-07-13 12:09:03 +08:00

157 lines
5.2 KiB
Python

"""Tests for PlanExecuteAgent: linear, replan, replan exhaustion, step budget, diffs."""
from __future__ import annotations
import os
import sys
import unittest
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.dirname(HERE))
from main import ( # noqa: E402
PlanDiff,
PlanExecuteAgent,
SessionResult,
Step,
ToolFailure,
make_deterministic_planner,
)
def perfect_executor(tool, args):
return f"ok:{tool}"
class TestLinear(unittest.TestCase):
def test_linear_plan_completes(self) -> None:
agent = PlanExecuteAgent(
planner=make_deterministic_planner(fail_step_id=None),
executor=perfect_executor,
)
res = agent.run("g")
self.assertEqual(res.status, "completed")
self.assertEqual(res.reason, "goal_met")
self.assertEqual(len(res.history), 4)
self.assertEqual([s.id for s in res.history], [1, 2, 3, 4])
self.assertEqual(res.revisions, [])
def test_no_plan_initial(self) -> None:
def empty(g, h, e):
return []
agent = PlanExecuteAgent(planner=empty, executor=perfect_executor)
res = agent.run("g")
self.assertEqual(res.status, "failed")
self.assertEqual(res.reason, "no_plan")
class TestReplan(unittest.TestCase):
def test_replan_once_on_transform_failure(self) -> None:
calls = {"transform_v1": 0, "transform_v2": 0}
def executor(tool, args):
if tool == "transform":
mode = args.get("mode")
if mode == "v1":
calls["transform_v1"] += 1
raise ToolFailure("transform v1 down")
if mode == "v2":
calls["transform_v2"] += 1
return "ok"
raise ToolFailure("transform unknown mode")
return f"ok:{tool}"
agent = PlanExecuteAgent(
planner=make_deterministic_planner(fail_step_id=2, recovery="route_around"),
executor=executor,
)
res = agent.run("g")
self.assertEqual(res.status, "completed")
self.assertEqual(res.reason, "goal_met")
self.assertEqual(calls["transform_v1"], 1)
self.assertEqual(calls["transform_v2"], 1)
self.assertEqual(len(res.revisions), 1)
def test_replan_diff_event_emitted(self) -> None:
def executor(tool, args):
if tool == "transform" and args.get("mode") == "v1":
raise ToolFailure("transform v1 boom")
return "ok"
agent = PlanExecuteAgent(
planner=make_deterministic_planner(None, recovery="route_around"),
executor=executor,
)
res = agent.run("g")
diff_events = [e for e in res.events if e.type == "plan.diff"]
self.assertEqual(len(diff_events), 1)
diff = diff_events[0].payload
self.assertEqual(diff["revision"], 1)
self.assertIn("removed", diff)
self.assertIn("added", diff)
self.assertIn("revised", diff)
self.assertIn(2, diff["revised"])
def test_replan_exhaustion_returns_failed(self) -> None:
def always_bad(tool, args):
if tool == "transform":
raise ToolFailure("nope")
return "ok"
def planner(g, h, e):
if e is None:
return [
Step(1, "fetch", {}, "fetch"),
Step(2, "transform", {"mode": "v1"}, "transform"),
]
return [
Step(2, "transform", {"mode": "v1"}, "transform again"),
]
agent = PlanExecuteAgent(planner=planner, executor=always_bad, max_replans=2, max_steps=50)
res = agent.run("g")
self.assertEqual(res.status, "failed")
self.assertEqual(res.reason, "replan_budget")
class TestBudgets(unittest.TestCase):
def test_step_budget_caps_execution(self) -> None:
def planner(g, h, e):
return [Step(i, "noop", {}, "noop") for i in range(20)]
agent = PlanExecuteAgent(planner=planner, executor=perfect_executor, max_steps=5)
res = agent.run("g")
self.assertEqual(res.status, "failed")
self.assertEqual(res.reason, "step_budget")
self.assertEqual(len(res.history), 5)
def test_max_replans_zero_returns_after_first_failure(self) -> None:
def planner(g, h, e):
return [Step(1, "bad", {}, "bad")]
def boom(tool, args):
raise ToolFailure("nope")
agent = PlanExecuteAgent(planner=planner, executor=boom, max_replans=0)
res = agent.run("g")
self.assertEqual(res.status, "failed")
self.assertEqual(res.reason, "replan_budget")
class TestEvents(unittest.TestCase):
def test_event_order(self) -> None:
agent = PlanExecuteAgent(
planner=make_deterministic_planner(None),
executor=perfect_executor,
)
res = agent.run("g")
types = [e.type for e in res.events]
self.assertEqual(types[0], "plan.commit")
self.assertEqual(types[-1], "session.complete")
self.assertEqual(types.count("step.start"), 4)
self.assertEqual(types.count("step.end"), 4)
if __name__ == "__main__":
unittest.main()