import unittest
from agent import calculator, lookup, run_agent, MockLLM

class Scripted:
    """Plays back fixed replies, so we can test the loop itself."""
    def __init__(self, replies): self.replies = list(replies)
    def complete(self, messages): return self.replies.pop(0)

class Tests(unittest.TestCase):
    def test_calculator(self):
        self.assertEqual(calculator("(12 + 8) * 3"), "60")
    def test_calculator_blocks_code(self):
        with self.assertRaises(Exception):
            calculator("__import__('os').system('echo hi')")
    def test_lookup_missing(self):
        with self.assertRaises(KeyError): lookup("banana")
    def test_mock_math(self):
        self.assertEqual(run_agent("What is (12 + 8) * 3?", MockLLM(), verbose=False), "60")
    def test_mock_lookup(self):
        out = run_agent("What is a workflow?", MockLLM(), verbose=False)
        self.assertIn("path the developer wrote", out)
    def test_unknown_tool_is_reported_not_crashed(self):
        llm = Scripted(['{"action": "email", "input": {}}', '{"final": "ok"}'])
        self.assertEqual(run_agent("hi", llm, verbose=False), "ok")
    def test_bad_json_retry(self):
        llm = Scripted(["not json", '{"final": "done"}'])
        self.assertEqual(run_agent("hi", llm, verbose=False), "done")
    def test_step_limit(self):
        llm = Scripted(['{"action": "calculator", "input": {"expression": "1+1"}}'] * 3)
        self.assertEqual(run_agent("hi", llm, max_steps=3, verbose=False),
                         "Stopped: step limit reached.")
    def test_tool_error_goes_back_to_model(self):
        llm = Scripted(['{"action": "calculator", "input": {"expression": "1/0"}}',
                        '{"final": "handled"}'])
        self.assertEqual(run_agent("hi", llm, verbose=False), "handled")

if __name__ == "__main__":
    unittest.main(verbosity=2)