tests/test_code_prompt_construction.py

"""Prompt construction for the `code` schema.

MBPP's gold tests call the reference function by name, but the MBPP `text`
field is a sentence of English that never names it. Prompting with the text
alone leaves the model to invent a name, so every candidate fails with
`NameError` and the item scores 0 for every model in every condition -- an
unscoreable item rather than a hard one. These tests pin the fix (the task's
test assertions travel with the prompt, as published MBPP evaluation does) and
pin the boundary: HumanEval prompts, which already are the function signature,
must keep their previous construction byte-for-byte.
"""
import unittest

from owmi.benchmarks.base import BenchmarkSpec
from owmi.benchmarks.evaluate import evaluate_code
from owmi.benchmarks.presets import DEFAULT_BENCHMARK_PRESETS
from owmi.benchmarks.schemas import example_from_row

SUFFIX = '\n\nWrite the function/code only.'

# Shaped exactly like the frozen MBPP rows: `text` is prose, `test_list` is a
# list of standalone asserts naming `rearange_string`, `code` is the reference.
MBPP_ROW = {
    'task_id': 39,
    'text': ('Write a function to check if the letters of a given string can be '
             'rearranged so that two characters that are adjacent to each other '
             'are different.'),
    'code': 'def rearange_string(S):\n    return S\n',
    'test_list': [
        'assert rearange_string("aab")==(\'aba\')',
        'assert rearange_string("aabb")==(\'abab\')',
        'assert rearange_string("abccdd")==(\'cdabcd\')',
    ],
    'test_setup_code': '',
    'challenge_test_list': [],
    'id': '39',
}

# Shaped like the frozen HumanEval rows: the prompt *is* the signature stub and
# the tests are one `check(candidate)` blob.
HUMANEVAL_ROW = {
    'task_id': 'HumanEval/0',
    'prompt': ('from typing import List\n\n\n'
               'def has_close_elements(numbers: List[float], threshold: float) -> bool:\n'
               '    """ Check if in given list of numbers, are any two numbers closer\n'
               '    than the given threshold.\n    """\n'),
    'canonical_solution': '    return False\n',
    'test': ('\n\nMETADATA = {}\n\n\ndef check(candidate):\n'
             '    assert candidate([1.0, 2.0, 3.9], 0.3) == True\n'),
    'entry_point': 'has_close_elements',
}


class CodePromptConstructionTests(unittest.TestCase):
    def test_mbpp_prompt_exposes_the_required_function_name(self):
        spec = DEFAULT_BENCHMARK_PRESETS['mbpp']
        example = example_from_row(spec, MBPP_ROW, 0)
        # The name the gold tests call must be discoverable from the prompt.
        self.assertIn('rearange_string', example.prompt)
        for assertion in MBPP_ROW['test_list']:
            self.assertIn(assertion, example.prompt)
        self.assertIn(MBPP_ROW['text'], example.prompt)
        self.assertTrue(example.prompt.endswith(SUFFIX))
        # The reference the scorer uses is untouched by the prompt change.
        self.assertEqual(example.reference['tests'], MBPP_ROW['test_list'])

    def test_mbpp_prompt_makes_the_item_winnable_end_to_end(self):
        """A model that reads the prompt can now produce a passing solution.

        Before the fix no candidate could pass, because the name in the tests
        was never disclosed. This closes the loop through the real executor.
        """
        spec = DEFAULT_BENCHMARK_PRESETS['mbpp']
        example = example_from_row(spec, MBPP_ROW, 0)
        answer = (
            '```python\n'
            'import heapq\n'
            'from collections import Counter\n'
            'def rearange_string(S):\n'
            '    ctr = Counter(S)\n'
            '    heap = [(-v, k) for k, v in ctr.items()]\n'
            '    heapq.heapify(heap)\n'
            '    if (-heap[0][0]) * 2 > len(S) + 1:\n'
            '        return ""\n'
            '    ans = []\n'
            '    while len(heap) >= 2:\n'
            '        n1, c1 = heapq.heappop(heap)\n'
            '        n2, c2 = heapq.heappop(heap)\n'
            '        ans.extend([c1, c2])\n'
            '        if n1 + 1: heapq.heappush(heap, (n1 + 1, c1))\n'
            '        if n2 + 1: heapq.heappush(heap, (n2 + 1, c2))\n'
            '    return "".join(ans) + (heap[0][1] if heap else "")\n'
            '```'
        )
        self.assertEqual(evaluate_code(answer, example, allow_exec=True), 1.0)
        # A candidate that invents its own name is what the old prompt forced,
        # and it still correctly scores 0 -- the fix removes the cause, not the
        # scorer's ability to fail a wrong answer.
        wrong_name = answer.replace('rearange_string', 'can_rearrange')
        self.assertEqual(evaluate_code(wrong_name, example, allow_exec=True), 0.0)

    def test_humaneval_prompt_construction_is_unchanged(self):
        spec = DEFAULT_BENCHMARK_PRESETS['humaneval']
        example = example_from_row(spec, HUMANEVAL_ROW, 0)
        self.assertEqual(example.prompt, HUMANEVAL_ROW['prompt'].rstrip() + SUFFIX)
        # The stub already names the function, so nothing is appended; in
        # particular the `check(candidate)` blob and its expected outputs must
        # not leak into a stub-completion prompt.
        self.assertNotIn('assert', example.prompt)
        self.assertNotIn('check(', example.prompt)
        self.assertNotIn('Your code should pass these tests', example.prompt)

    def test_code_item_without_tests_still_builds_a_prompt(self):
        spec = BenchmarkSpec('toy', 'code', 'unused', field_map={
            'question': 'text', 'solution': 'code', 'tests': 'test_list',
        })
        row = {'id': '1', 'text': 'Write a function that adds two numbers.', 'code': 'x'}
        example = example_from_row(spec, row, 0)
        self.assertEqual(example.prompt, row['text'] + SUFFIX)
        self.assertNotIn('Your code should pass these tests', example.prompt)

        # An empty-but-present tests field is the same case.
        empty = example_from_row(spec, {**row, 'test_list': []}, 0)
        self.assertEqual(empty.prompt, row['text'] + SUFFIX)
        blank = example_from_row(spec, {**row, 'test_list': ['', '  ']}, 0)
        self.assertEqual(blank.prompt, row['text'] + SUFFIX)

    def test_include_tests_in_prompt_overrides_the_signature_heuristic(self):
        forced_off = BenchmarkSpec('mbpp_off', 'code', 'unused', field_map={
            'question': 'text', 'solution': 'code', 'tests': 'test_list',
            'include_tests_in_prompt': False,
        })
        example = example_from_row(forced_off, MBPP_ROW, 0)
        self.assertNotIn('rearange_string', example.prompt)

        forced_on = BenchmarkSpec('humaneval_on', 'code', 'unused', field_map={
            'question': 'prompt', 'solution': 'canonical_solution', 'tests': 'test',
            'entry_point': 'entry_point', 'include_tests_in_prompt': True,
        })
        example = example_from_row(forced_on, HUMANEVAL_ROW, 0)
        self.assertIn('def check(candidate)', example.prompt)

    def test_tests_are_rendered_for_both_list_and_string_shapes(self):
        list_spec = BenchmarkSpec('list_tests', 'code', 'unused', field_map={
            'question': 'text', 'tests': 'test_list', 'include_tests_in_prompt': True,
        })
        joined = example_from_row(list_spec, MBPP_ROW, 0).prompt
        self.assertIn('assert rearange_string("aab")', joined)
        self.assertIn('assert rearange_string("abccdd")', joined)

        str_spec = BenchmarkSpec('str_tests', 'code', 'unused', field_map={
            'question': 'text', 'tests': 'test', 'include_tests_in_prompt': True,
        })
        row = {'id': '2', 'text': 'Add two numbers.', 'test': 'assert add(1, 2) == 3'}
        self.assertIn('assert add(1, 2) == 3', example_from_row(str_spec, row, 0).prompt)


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