tests/unit/test_agent_runner.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from agent.runner import (
    AGENT_SYSTEM_PROMPT,
    default_safety_gates,
    export_agent_traces,
    export_agent_traces_hf_dataset,
    run_agent_loop,
    run_paper_to_code_loop,
    save_agent_trace,
)


class AgentRunnerTest(unittest.TestCase):
    def test_agent_loop_creates_research_plan_implementation_trace(self) -> None:
        session = run_agent_loop("Improve dataset stats")

        phases = [step.phase for step in session.steps]
        self.assertIn("research", phases)
        self.assertIn("plan", phases)
        self.assertIn("implement", phases)
        self.assertIn("verify", phases)
        self.assertIn("safe_calculator", session.tools)
        self.assertTrue(session.safety_gates)
        self.assertEqual(session.as_dict()["system_prompt"], AGENT_SYSTEM_PROMPT)

    def test_agent_loop_can_use_safe_calculator_tool(self) -> None:
        session = run_agent_loop("calculate: 2 + 2")

        self.assertIn("tool:safe_calculator", [step.phase for step in session.steps])
        self.assertIn("4.0", session.as_markdown())

    def test_saves_and_exports_agent_trace(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            source = Path(tmp) / "agent_traces.jsonl"
            output = Path(tmp) / "exports" / "agent_traces.jsonl"
            session = run_agent_loop("Improve docs")

            saved = save_agent_trace(session, source)
            exported = export_agent_traces(source, output)

            self.assertEqual(saved, source)
            self.assertEqual(exported, output)
            self.assertIn("Improve docs", output.read_text(encoding="utf-8"))

    def test_exports_agent_trace_dataset(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            source = Path(tmp) / "agent_traces.jsonl"
            target = Path(tmp) / "dataset"
            save_agent_trace(run_agent_loop("Trace me"), source)

            exported = export_agent_traces_hf_dataset(source, target)

            self.assertEqual(exported, target)
            self.assertTrue((target / "data.jsonl").exists())
            self.assertTrue((target / "README.md").exists())

    def test_paper_to_code_trace_has_required_phases_and_safety_gates(self) -> None:
        session = run_paper_to_code_loop(
            "Demo paper",
            "Claims a local reward model can rank completions.",
            "Implement deterministic reward eval.",
        )

        phases = [step.phase for step in session.steps]
        self.assertEqual(phases, ["research", "plan", "implement", "verify"])
        self.assertIn("Paper-to-code", session.task)
        self.assertIn("No shell commands", session.as_markdown())

    def test_default_safety_gates_block_external_side_effects(self) -> None:
        gates = default_safety_gates()

        self.assertTrue(any("downloaded automatically" in gate for gate in gates))
        self.assertTrue(any("matching test" in gate for gate in gates))


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