See More

"""Tests for the FlyPython deterministic troubleshooting agent.""" import threading import unittest from unittest.mock import patch from agent_loop import Action, PLAYBOOKS, TOOLS, Task, classify_issue, run_agent class AgentLoopTests(unittest.TestCase): def test_network_issue_returns_playbook(self) -> None: task = run_agent(Task("My API request times out")) self.assertEqual(task.observations[0], "network") self.assertEqual(task.answer, PLAYBOOKS["network"]) self.assertEqual( task.trace, ["1. classify_issue", "2. lookup_playbook", "3. finish"], ) def test_dependency_issue_returns_playbook(self) -> None: task = run_agent(Task("An imported package is missing")) self.assertEqual(task.observations[0], "dependency") self.assertEqual(task.answer, PLAYBOOKS["dependency"]) def test_general_issue_returns_playbook(self) -> None: task = run_agent(Task("Capitalization issue in a generated title")) self.assertEqual(classify_issue("Capitalization issue"), "general") self.assertEqual(task.observations[0], "general") self.assertEqual(task.answer, PLAYBOOKS["general"]) def test_step_budget_stops_the_loop(self) -> None: with self.assertRaisesRegex(RuntimeError, "1-step budget"): run_agent(Task("An import fails"), max_steps=1) def test_unknown_tool_is_rejected(self) -> None: def unsafe_planner(_task: Task) -> Action: return Action("delete_files") with self.assertRaisesRegex(ValueError, "Unknown tool"): run_agent(Task("Delete everything"), planner=unsafe_planner) def test_non_action_planner_result_is_rejected(self) -> None: def invalid_planner(_task: Task) -> object: return {"name": "finish", "arguments": {"answer": "unsafe"}} with self.assertRaisesRegex(TypeError, "Planner must return an Action"): run_agent(Task("Return an invalid result"), planner=invalid_planner) def test_invalid_action_fields_are_rejected(self) -> None: invalid_actions = ( (Action(123, {}), "Action name must be a string"), (Action("finish", []), "Action arguments must be a dictionary"), ) for action, error_message in invalid_actions: with self.subTest(error_message=error_message): with self.assertRaisesRegex(TypeError, error_message): run_agent(Task("Reject malformed action"), planner=lambda _: action) def test_invalid_argument_type_is_rejected(self) -> None: def invalid_planner(_task: Task) -> Action: return Action("classify_issue", {"text": 123}) with self.assertRaisesRegex( TypeError, "Argument 'text' for classify_issue must be str" ): run_agent(Task("Use an invalid argument"), planner=invalid_planner) def test_invalid_finish_answer_type_is_rejected(self) -> None: def invalid_planner(_task: Task) -> Action: return Action("finish", {"answer": 123}) with self.assertRaisesRegex( TypeError, "Argument 'answer' for finish must be str" ): run_agent(Task("Finish unsafely"), planner=invalid_planner) def test_finish_rejects_extra_arguments(self) -> None: def invalid_planner(_task: Task) -> Action: return Action("finish", {"answer": "done", "extra": "unsafe"}) with self.assertRaisesRegex(ValueError, "Invalid arguments for finish"): run_agent(Task("Finish with extra data"), planner=invalid_planner) def test_planner_timeout_returns_control_from_blocking_call(self) -> None: release = threading.Event() def blocking_planner(_task: Task) -> Action: release.wait() return Action("finish", {"answer": "too late"}) try: with self.assertRaisesRegex( TimeoutError, "Planner exceeded its 0.01-second timeout" ): run_agent( Task("Wait forever"), planner=blocking_planner, deadline_seconds=1.0, planner_timeout_seconds=0.01, ) finally: release.set() def test_tool_timeout_returns_control_from_blocking_call(self) -> None: release = threading.Event() def blocking_tool(text: str) -> str: del text release.wait() return "general" try: with patch.dict(TOOLS, {"classify_issue": blocking_tool}): with self.assertRaisesRegex( TimeoutError, "Tool classify_issue exceeded its 0.01-second timeout", ): run_agent( Task("Wait forever"), deadline_seconds=1.0, tool_timeout_seconds=0.01, ) finally: release.set() def test_wall_clock_deadline_caps_a_blocking_planner(self) -> None: release = threading.Event() def blocking_planner(_task: Task) -> Action: release.wait() return Action("finish", {"answer": "too late"}) try: with self.assertRaisesRegex( TimeoutError, "0.01-second wall-clock deadline" ): run_agent( Task("Wait forever"), planner=blocking_planner, deadline_seconds=0.01, planner_timeout_seconds=1.0, ) finally: release.set() if __name__ == "__main__": unittest.main()