PocketAccountant: custom ledger UI + deterministic agent (engine, ledger, retrieval, classifier)
c55ab5e verified | """Tests for the typed tools and the agent loop (no model — ScriptedClient). | |
| Run from the project root: python -m unittest discover -s tests | |
| """ | |
| import json | |
| import unittest | |
| from src.agent import ( | |
| Agent, | |
| AssistantTurn, | |
| ScriptedClient, | |
| ToolCall, | |
| ToolContext, | |
| build_tools, | |
| ) | |
| from src.ledger import Ledger | |
| def seeded_ledger(): | |
| lg = Ledger(":memory:") | |
| lg.record_income("2024-05-10", "Diseño", 45000, iva_rate="0.16") | |
| lg.record_expense("2024-05-15", "Adobe", 12000, iva_rate="0.16") | |
| return lg | |
| class TestTools(unittest.TestCase): | |
| def setUp(self): | |
| self.ctx = ToolContext(seeded_ledger()) | |
| self.tools = build_tools(self.ctx) | |
| def tearDown(self): | |
| self.ctx.ledger.close() | |
| def test_record_then_summary(self): | |
| self.tools["record_income"].handler( | |
| date="2024-05-25", description="Extra", amount=5000, iva_rate=0.16) | |
| summary = self.tools["month_summary"].handler(year=2024, month=5) | |
| self.assertEqual(summary["income"], "50000.00") # 45,000 + 5,000 | |
| def test_compute_isr_resico(self): | |
| r = self.tools["compute_isr_resico"].handler(year=2024, month=5) | |
| self.assertEqual(r["amount"], "495.00") | |
| self.assertIn("113-E", r["source"]) # provenance carried through | |
| def test_compare_regimes(self): | |
| r = self.tools["compare_regimes"].handler(year=2024, month=5) | |
| self.assertEqual(r["recommended"], "RESICO") | |
| self.assertEqual(r["resico_isr"], "495.00") | |
| self.assertEqual(r["general_isr"], "5418.90") | |
| def test_us_tax_estimate(self): | |
| r = self.tools["us_tax_estimate"].handler(gross_income=100000, expenses=40000) | |
| self.assertEqual(r["net_profit"]["amount"], "60000.00") | |
| self.assertEqual(r["self_employment_tax"]["amount"], "8477.73") | |
| def test_every_tool_has_schema(self): | |
| for name, tool in self.tools.items(): | |
| s = tool.schema() | |
| self.assertEqual(s["function"]["name"], name) | |
| self.assertIn("parameters", s["function"]) | |
| class TestAgentLoop(unittest.TestCase): | |
| def setUp(self): | |
| self.ctx = ToolContext(seeded_ledger()) | |
| self.tools = build_tools(self.ctx) | |
| def tearDown(self): | |
| self.ctx.ledger.close() | |
| def test_two_step_run(self): | |
| client = ScriptedClient([ | |
| AssistantTurn(tool_calls=[ToolCall("compare_regimes", {"year": 2024, "month": 5})]), | |
| AssistantTurn(text="En mayo te conviene RESICO: pagas 495 vs 5,418.90."), | |
| ]) | |
| agent = Agent(client, self.tools) | |
| trace = agent.run("¿Qué régimen me conviene en mayo?") | |
| self.assertEqual(len(trace.steps), 1) | |
| self.assertEqual(trace.steps[0].tool, "compare_regimes") | |
| self.assertTrue(trace.steps[0].ok) | |
| self.assertEqual(trace.steps[0].result["recommended"], "RESICO") | |
| self.assertIn("RESICO", trace.final_answer) | |
| def test_tool_error_is_surfaced_not_raised(self): | |
| # Missing required 'month' → handler raises → loop records ok=False, continues. | |
| client = ScriptedClient([ | |
| AssistantTurn(tool_calls=[ToolCall("compute_isr_resico", {"year": 2024})]), | |
| AssistantTurn(text="Necesito el mes para calcular tu ISR."), | |
| ]) | |
| agent = Agent(client, self.tools) | |
| trace = agent.run("¿Cuánto ISR debo?") | |
| self.assertFalse(trace.steps[0].ok) | |
| self.assertIn("error", trace.steps[0].result) | |
| self.assertIn("mes", trace.final_answer) | |
| def test_unknown_tool(self): | |
| client = ScriptedClient([ | |
| AssistantTurn(tool_calls=[ToolCall("does_not_exist", {})]), | |
| AssistantTurn(text="done"), | |
| ]) | |
| trace = Agent(client, self.tools).run("hi") | |
| self.assertFalse(trace.steps[0].ok) | |
| self.assertIn("unknown tool", trace.steps[0].result["error"]) | |
| def test_trace_serializes(self): | |
| client = ScriptedClient([ | |
| AssistantTurn(tool_calls=[ToolCall("compute_iva", {"year": 2024, "month": 5})]), | |
| AssistantTurn(text="Tu IVA a cargo es 5,280.00."), | |
| ]) | |
| trace = Agent(client, self.tools).run("¿Cuánto IVA debo en mayo?") | |
| blob = trace.to_json() | |
| parsed = json.loads(blob) | |
| self.assertEqual(parsed["steps"][0]["tool"], "compute_iva") | |
| self.assertEqual(parsed["steps"][0]["result"]["amount"], "5280.00") | |
| def test_step_budget_terminates(self): | |
| # Client always asks for a tool call; the loop must stop at max_steps. | |
| loop_turn = AssistantTurn(tool_calls=[ToolCall("month_summary", {"year": 2024, "month": 5})]) | |
| client = ScriptedClient([loop_turn] * 20) | |
| trace = Agent(client, self.tools, max_steps=3).run("loop forever") | |
| self.assertEqual(len(trace.steps), 3) | |
| self.assertIn("step limit", trace.final_answer) | |
| if __name__ == "__main__": | |
| unittest.main() | |