Spaces:
Running
Running
File size: 6,440 Bytes
1828bcf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """
Unit tests for scripts/run_agent.py and scripts/build_index.py CLIs.
We mock out heavy dependencies (LangGraphAgent, MathLibRetriever) so the tests
exercise only the argparse / control-flow surface of each script.
"""
import contextlib
import importlib
import io
import os
import sys
import unittest
from unittest import mock
_HERE = os.path.dirname(__file__)
_ROOT = os.path.abspath(os.path.join(_HERE, ".."))
sys.path.insert(0, os.path.join(_ROOT, "scripts"))
sys.path.insert(0, os.path.join(_ROOT, "src"))
def _fresh_import(module_name: str):
"""(Re)import a script module so that patches applied via sys.modules
take effect for top-level `from x import Y` statements."""
if module_name in sys.modules:
del sys.modules[module_name]
return importlib.import_module(module_name)
# ---------------------------------------------------------------------------
# run_agent.py
# ---------------------------------------------------------------------------
class TestRunAgent(unittest.TestCase):
def _run_main(self, argv, agent_success=True):
"""Patch LangGraphAgent + sys.argv, import run_agent, run main(),
return (captured_stdout, MockAgent class, mock_instance)."""
mock_agent_cls = mock.MagicMock(name="LangGraphAgent")
mock_instance = mock_agent_cls.return_value
mock_instance.solve_file.return_value = agent_success
# Inject the mock into the source module so `from langgraph_agent import LangGraphAgent`
# picks it up at import-time inside run_agent.
with mock.patch.dict(sys.modules, {"langgraph_agent": mock.MagicMock(LangGraphAgent=mock_agent_cls)}):
run_agent = _fresh_import("run_agent")
buf = io.StringIO()
with mock.patch.object(sys, "argv", argv), contextlib.redirect_stdout(buf):
run_agent.main()
return buf.getvalue(), mock_agent_cls, mock_instance
def test_defaults(self):
out, cls, inst = self._run_main(["run_agent.py", "problems/simple_add.lean"])
cls.assert_called_once_with(
model_name="qwen3-vl:4b",
max_retries=5,
index_dir=None,
)
inst.solve_file.assert_called_once_with("problems/simple_add.lean")
self.assertIn("Success", out)
def test_custom_model(self):
_, cls, _ = self._run_main(
["run_agent.py", "problems/simple_add.lean", "--model", "custom-model"]
)
_, kwargs = cls.call_args
self.assertEqual(kwargs["model_name"], "custom-model")
def test_retries_forwarded(self):
_, cls, _ = self._run_main(
["run_agent.py", "problems/simple_add.lean", "--retries", "3"]
)
_, kwargs = cls.call_args
self.assertEqual(kwargs["max_retries"], 3)
def test_index_dir_forwarded(self):
_, cls, _ = self._run_main(
["run_agent.py", "problems/simple_add.lean", "--index-dir", "/tmp/idx"]
)
_, kwargs = cls.call_args
self.assertEqual(kwargs["index_dir"], "/tmp/idx")
def test_failure_prints_failed_message(self):
out, _, _ = self._run_main(
["run_agent.py", "problems/simple_add.lean"], agent_success=False
)
self.assertIn("Failed", out)
self.assertNotIn("Success", out)
def test_missing_file_arg_exits(self):
mock_agent_cls = mock.MagicMock(name="LangGraphAgent")
with mock.patch.dict(sys.modules, {"langgraph_agent": mock.MagicMock(LangGraphAgent=mock_agent_cls)}):
run_agent = _fresh_import("run_agent")
with mock.patch.object(sys, "argv", ["run_agent.py"]):
# argparse writes the error to stderr; swallow it for clean test output.
with contextlib.redirect_stderr(io.StringIO()):
with self.assertRaises(SystemExit) as ctx:
run_agent.main()
self.assertEqual(ctx.exception.code, 2)
# ---------------------------------------------------------------------------
# build_index.py
# ---------------------------------------------------------------------------
class TestBuildIndex(unittest.TestCase):
def _run_main(self, argv, is_built=False):
"""Patch MathLibRetriever + sys.argv, import build_index, run main(),
return (captured_stdout, MockRetriever class, mock_instance)."""
mock_retr_cls = mock.MagicMock(name="MathLibRetriever")
mock_instance = mock_retr_cls.return_value
mock_instance.is_index_built.return_value = is_built
mock_instance.index_dir = "/fake/index/dir"
with mock.patch.dict(sys.modules, {"retriever": mock.MagicMock(MathLibRetriever=mock_retr_cls)}):
build_index = _fresh_import("build_index")
buf = io.StringIO()
with mock.patch.object(sys, "argv", argv), contextlib.redirect_stdout(buf):
build_index.main()
return buf.getvalue(), mock_retr_cls, mock_instance
def test_defaults(self):
out, cls, inst = self._run_main(["build_index.py"])
cls.assert_called_once_with(index_dir=None)
inst.build.assert_called_once_with(mathlib_root=None, max_files=None)
self.assertIn("Done", out)
def test_mathlib_root_forwarded(self):
_, _, inst = self._run_main(["build_index.py", "--mathlib-root", "/tmp/fake"])
_, kwargs = inst.build.call_args
self.assertEqual(kwargs["mathlib_root"], "/tmp/fake")
def test_max_files_forwarded(self):
_, _, inst = self._run_main(["build_index.py", "--max-files", "5"])
_, kwargs = inst.build.call_args
self.assertEqual(kwargs["max_files"], 5)
def test_index_dir_forwarded(self):
_, cls, _ = self._run_main(["build_index.py", "--index-dir", "/tmp/idx"])
cls.assert_called_once_with(index_dir="/tmp/idx")
def test_existing_index_skips_rebuild(self):
out, _, inst = self._run_main(["build_index.py"], is_built=True)
inst.build.assert_not_called()
self.assertIn("Index already exists", out)
def test_existing_index_rebuilds_when_max_files_set(self):
out, _, inst = self._run_main(
["build_index.py", "--max-files", "5"], is_built=True
)
inst.build.assert_called_once_with(mathlib_root=None, max_files=5)
self.assertNotIn("Index already exists", out)
if __name__ == "__main__":
unittest.main()
|