"""Unit tests for pure helpers in scripts/benchmark.py. Covered helpers: - load_local_problems(problems_dir, max_problems) - pass_at_k(results, k) - _ensure_import_and_sorry(code) """ import os import sys import tempfile import unittest from pathlib import Path from unittest import mock sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "scripts"))) sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src"))) # Stub out the heavy `langgraph_agent` import so `benchmark` loads without # its full dependency chain (langchain_anthropic, etc.). sys.modules.setdefault( "langgraph_agent", mock.MagicMock(LangGraphAgent=mock.MagicMock()) ) from benchmark import _ensure_import_and_sorry, load_local_problems, pass_at_k # noqa: E402 # --------------------------------------------------------------------------- # load_local_problems # --------------------------------------------------------------------------- class LoadLocalProblemsTests(unittest.TestCase): def test_filters_files_without_sorry(self): """A complete (no sorry) file is skipped; a withsorry file is loaded.""" with tempfile.TemporaryDirectory() as td: root = Path(td) (root / "complete.lean").write_text( "import Mathlib\n\ntheorem foo : 1 + 1 = 2 := by rfl\n", encoding="utf-8", ) (root / "withsorry.lean").write_text( "import Mathlib\n\ntheorem bar : 1 + 1 = 2 := by sorry\n", encoding="utf-8", ) rows = load_local_problems(td) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["name"], "withsorry") def test_max_problems_limits_loading(self): """max_problems=2 with 5 sorry-files yields exactly the first 2 sorted.""" with tempfile.TemporaryDirectory() as td: root = Path(td) names = ["e_fifth", "a_first", "b_second", "c_third", "d_fourth"] for n in names: (root / f"{n}.lean").write_text( f"-- {n}\nimport Mathlib\nexample : True := by sorry\n", encoding="utf-8", ) rows = load_local_problems(td, max_problems=2) self.assertEqual(len(rows), 2) # Sorted alphabetically: a_first, b_second are the first 2 self.assertEqual(rows[0]["name"], "a_first") self.assertEqual(rows[1]["name"], "b_second") def test_max_problems_none_loads_all(self): """max_problems=None loads all sorry-files.""" with tempfile.TemporaryDirectory() as td: root = Path(td) for n in ["p1", "p2", "p3"]: (root / f"{n}.lean").write_text( "import Mathlib\nexample : True := by sorry\n", encoding="utf-8", ) rows = load_local_problems(td, max_problems=None) self.assertEqual(len(rows), 3) def test_missing_dir_raises_runtime_error(self): """Non-existent directory raises RuntimeError with a clear message.""" missing = "/tmp/definitely_does_not_exist_xyz_12345abc" with self.assertRaises(RuntimeError) as ctx: load_local_problems(missing) self.assertIn(missing, str(ctx.exception)) def test_empty_dir_returns_empty_list(self): """An empty directory returns [] without crashing.""" with tempfile.TemporaryDirectory() as td: rows = load_local_problems(td) self.assertEqual(rows, []) def test_name_is_file_stem(self): """`name` is the file stem (no .lean suffix).""" with tempfile.TemporaryDirectory() as td: root = Path(td) (root / "lecture4_1.lean").write_text( "import Mathlib\nexample : True := by sorry\n", encoding="utf-8", ) rows = load_local_problems(td) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["name"], "lecture4_1") def test_ignores_non_lean_files(self): """Non-.lean files (e.g. notes.txt) are ignored even if they contain 'sorry'.""" with tempfile.TemporaryDirectory() as td: root = Path(td) (root / "notes.txt").write_text( "I'm sorry to say this is not Lean.\n", encoding="utf-8" ) (root / "prob.lean").write_text( "import Mathlib\nexample : True := by sorry\n", encoding="utf-8", ) rows = load_local_problems(td) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["name"], "prob") # --------------------------------------------------------------------------- # pass_at_k # --------------------------------------------------------------------------- class PassAtKTests(unittest.TestCase): def test_empty_results(self): self.assertEqual(pass_at_k([], 1), 0.0) self.assertEqual(pass_at_k([], 5), 0.0) def test_all_success_attempt_one(self): results = [ {"success": True, "solved_at_attempt": 1} for _ in range(5) ] self.assertEqual(pass_at_k(results, 1), 1.0) def test_three_success_at_attempt_two_plus_two_fail(self): results = ( [{"success": True, "solved_at_attempt": 2} for _ in range(3)] + [{"success": False, "solved_at_attempt": None} for _ in range(2)] ) self.assertAlmostEqual(pass_at_k(results, 2), 0.6) self.assertEqual(pass_at_k(results, 1), 0.0) def test_success_at_attempt_three_k_two(self): results = [{"success": True, "solved_at_attempt": 3}] self.assertEqual(pass_at_k(results, 2), 0.0) def test_success_at_attempt_three_k_ten(self): results = [ {"success": True, "solved_at_attempt": 3}, {"success": False, "solved_at_attempt": None}, {"success": False, "solved_at_attempt": None}, {"success": False, "solved_at_attempt": None}, ] # 1 solved out of 4 total self.assertEqual(pass_at_k(results, 10), 0.25) def test_all_failures(self): results = [ {"success": False, "solved_at_attempt": None} for _ in range(4) ] for k in (1, 2, 5, 100): self.assertEqual(pass_at_k(results, k), 0.0) # --------------------------------------------------------------------------- # _ensure_import_and_sorry # --------------------------------------------------------------------------- class EnsureImportAndSorryTests(unittest.TestCase): def test_missing_import_prepends_mathlib(self): code = "theorem foo : 1 + 1 = 2 := by sorry" out = _ensure_import_and_sorry(code) self.assertTrue(out.startswith("import Mathlib\n\n")) def test_has_import_and_sorry_unchanged(self): code = "import Mathlib\n\ntheorem foo : 1 + 1 = 2 := by sorry\n" out = _ensure_import_and_sorry(code) self.assertEqual(out, code) def test_has_import_missing_sorry_ends_by(self): code = "import Mathlib\n\ntheorem foo : 1 + 1 = 2 := by" out = _ensure_import_and_sorry(code) self.assertIn("sorry", out) # Should append sorry after "by" self.assertTrue(out.rstrip().endswith("sorry")) def test_has_import_missing_sorry_no_by(self): code = "import Mathlib\n\ntheorem foo : 1 + 1 = 2 :=" out = _ensure_import_and_sorry(code) self.assertIn("sorry", out) self.assertIn("by", out) # Should append " by\n sorry\n" self.assertTrue(out.rstrip().endswith("sorry")) def test_empty_string_no_crash_locks_current_behavior(self): """Empty input must not crash. Lock in observed output.""" out = _ensure_import_and_sorry("") # No import in "" → prepended. No ":= by" / ":=" → nothing else appended. self.assertEqual(out, "import Mathlib\n\n") if __name__ == "__main__": unittest.main()