diff --git a/engine/hooks/llm-judge/judge_test_base.py b/engine/hooks/llm-judge/judge_test_base.py new file mode 100644 index 00000000..01249d18 --- /dev/null +++ b/engine/hooks/llm-judge/judge_test_base.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import json +import os +import tempfile +import unittest +from unittest.mock import patch + +import judge + + +class JudgeTestCase(unittest.TestCase): + def setUp(self): + super().setUp() + self.state = tempfile.TemporaryDirectory() + self.judge_env = patch.dict(os.environ, {judge.STATE_ENV: self.state.name}) + self.judge_env.start() + os.environ.pop(judge.CHILD_ENV, None) + os.environ.pop(judge.RUNNERS_ENV, None) + + def tearDown(self): + self.judge_env.stop() + self.state.cleanup() + super().tearDown() + + def use_runners(self, *entries): + os.environ[judge.RUNNERS_ENV] = json.dumps(list(entries)) diff --git a/engine/hooks/llm-judge/tests/test_judge.py b/engine/hooks/llm-judge/tests/test_judge.py index 08d61c0d..c186e8ee 100644 --- a/engine/hooks/llm-judge/tests/test_judge.py +++ b/engine/hooks/llm-judge/tests/test_judge.py @@ -6,7 +6,6 @@ import json import os import sys -import tempfile import time import unittest import warnings @@ -16,6 +15,7 @@ sys.path.insert(0, LIB_DIR) import judge # noqa: E402 +from judge_test_base import JudgeTestCase # noqa: E402 PY = sys.executable @@ -31,21 +31,7 @@ def runner(name, script): SLOW_MATCH = runner("slow", "import json, sys, time; time.sleep(2); print(json.dumps({'match': True, 'prompt': sys.argv[1]}))") -class JudgeTestCase(unittest.TestCase): - def setUp(self): - self.state = tempfile.TemporaryDirectory() - self.env = patch.dict(os.environ, {judge.STATE_ENV: self.state.name}) - self.env.start() - os.environ.pop(judge.CHILD_ENV, None) - os.environ.pop(judge.RUNNERS_ENV, None) - - def tearDown(self): - self.env.stop() - self.state.cleanup() - - def use_runners(self, *entries): - os.environ[judge.RUNNERS_ENV] = json.dumps(list(entries)) - +class JudgeBehaviorTestCase(JudgeTestCase): def job(self, **overrides): base = { "id": "job-1", @@ -59,7 +45,7 @@ def job(self, **overrides): return base -class TestAsk(JudgeTestCase): +class TestAsk(JudgeBehaviorTestCase): def test_first_runner_fails_second_answers_and_one_failed_attempt_is_recorded(self): self.use_runners(EXIT_NONZERO, ANSWER_MATCH) result = judge.ask("hello judge") @@ -125,7 +111,7 @@ def test_default_runner_order_is_codex_then_claude_then_cursor(self): self.assertEqual([name for name, _ in judge.runners()], ["codex", "claude", "cursor"]) -class TestVerdict(JudgeTestCase): +class TestVerdict(JudgeBehaviorTestCase): def test_hit_when_every_hit_key_is_true(self): job = self.job(hit_if_all_true=["match", "sure"]) result = judge.verdict(job, {"outcome": "answered", "runner": "answers", "answer": {"match": True, "sure": True}, "attempts": []}) @@ -148,7 +134,7 @@ def test_unchecked_when_ask_was_unchecked(self): self.assertEqual(result["attempts"], attempts) -class TestBackground(JudgeTestCase): +class TestBackground(JudgeBehaviorTestCase): def test_enqueue_as_judge_child_returns_none_and_starts_nothing(self): os.environ[judge.CHILD_ENV] = "1" self.use_runners(ANSWER_MATCH)