tests/test_code_extraction.py

"""Regression tests for code extraction from model generations.

The fixtures marked OBSERVED are verbatim copies of persisted OWMI generations
(``<results-dir>/...__{humaneval,mbpp}__*/**/{baseline,intervention}_answer.txt``,
Qwen2.5-7B-Instruct, 2026-08-09). The suite-wide system prompt tells the model
to "Return only the requested JSON object", so code answers arrive wrapped in a
```json fence instead of a ```python one.
"""

from __future__ import annotations

import unittest

from owmi.benchmarks.base import BenchmarkExample
from owmi.benchmarks.evaluate import _extract_code_block, evaluate_code


# OBSERVED: humaneval/HumanEval_109, every condition except random_direction.
OBSERVED_HUMANEVAL_JSON = (
    '```json\n'
    '{\n'
    '  "function": "def move_one_ball(arr): \\n    n = len(arr)\\n    if n == 0:\\n'
    '        return True\\n    count = 0\\n    for i in range(n-1):\\n'
    '        if arr[i] > arr[i+1]:\\n            count += 1\\n'
    '            if count > 1:\\n                return False\\n'
    '    if arr[n-1] > arr[0]:\\n        count += 1\\n    return count <= 1"\n'
    '}\n'
    '```'
)

# OBSERVED: humaneval/HumanEval_109, condition_random_direction intervention.
# The perturbation knocked the model out of JSON mode; this path already worked.
OBSERVED_HUMANEVAL_PYTHON_FENCE = (
    '```python\n'
    'def move_one_ball(arr):\n'
    '    if len(arr) == 0:\n'
    '        return True\n'
    '    \n'
    '    count = 0\n'
    '    for i in range(len(arr)):\n'
    '        if arr[i] < arr[(i-1) % len(arr)]:\n'
    '            count += 1\n'
    '            if count > 1:\n'
    '                return False\n'
    '    return True\n'
    '```'
)

# OBSERVED: mbpp/39, condition_random_direction intervention. Literal newlines
# inside the JSON string make this blob unparseable by json.loads.
OBSERVED_MBPP_RAW_NEWLINE_JSON = '''```json
{
    "function_code": "
def can_rearrange(s):
    if len(s) < 2:
        return True
    stack = []
    for char in s:
        if stack and stack[-1] == char:
            return False
        stack.append(char)
    return True
"
}
```'''

# OBSERVED: mbpp/39 baseline, truncated by the 128-token budget mid-string:
# no closing quote, no closing brace, no closing fence.
OBSERVED_MBPP_TRUNCATED_JSON = (
    '```json\n'
    '{\n'
    '  "function": "def can_rearrange(s: str) -> bool:\\n    if len(s) <= 1:\\n'
    '        return True\\n    char_count = {}\\n    for char in s:\\n'
    '        if char in char_count:\\n            char_count[char] += 1\\n'
    '        else:\\n            char_count[char] = 1\\n'
    '    max_char_count = max(char_count.values())\\n'
    '    most_common_char = max(char_count, key=char_count.get)\\n'
    '    if max_char_count > (len(s) + 1) // 2:\\n        return'
)


class LegacyExtractionPathsUnchanged(unittest.TestCase):
    """The paths that worked before the JSON fix must behave identically."""

    def test_bare_fence(self):
        self.assertEqual(
            _extract_code_block('```\ndef add(a, b):\n    return a + b\n```'),
            'def add(a, b):\n    return a + b',
        )

    def test_python_fence(self):
        self.assertEqual(
            _extract_code_block('```python\ndef add(a, b):\n    return a + b\n```'),
            'def add(a, b):\n    return a + b',
        )

    def test_observed_python_fence(self):
        extracted = _extract_code_block(OBSERVED_HUMANEVAL_PYTHON_FENCE)
        self.assertTrue(extracted.startswith('def move_one_ball(arr):'))
        self.assertTrue(extracted.endswith('return True'))
        self.assertNotIn('```', extracted)

    def test_raw_unfenced_code(self):
        self.assertEqual(
            _extract_code_block('def add(a, b):\n    return a + b\n'),
            'def add(a, b):\n    return a + b',
        )

    def test_prose_around_python_fence_still_wins(self):
        text = 'Here you go:\n```python\ndef f():\n    return 1\n```\nHope that helps.'
        self.assertEqual(_extract_code_block(text), 'def f():\n    return 1')

    def test_python_fence_wins_over_a_json_fence_in_the_same_response(self):
        text = (
            '```json\n{"note": "def not_the_answer(): return 0"}\n```\n'
            '```python\ndef real(): return 1\n```'
        )
        self.assertEqual(_extract_code_block(text), 'def real(): return 1')


class JsonWrappedExtraction(unittest.TestCase):
    """The shapes actually observed in persisted generations."""

    def test_observed_humaneval_function_key(self):
        extracted = _extract_code_block(OBSERVED_HUMANEVAL_JSON)
        self.assertTrue(extracted.startswith('def move_one_ball(arr):'))
        self.assertIn('return count <= 1', extracted)
        self.assertNotIn('"function"', extracted)
        self.assertNotIn('\\n', extracted)
        compile(extracted, '<t>', 'exec')

    def test_observed_mbpp_function_code_key_with_literal_newlines(self):
        extracted = _extract_code_block(OBSERVED_MBPP_RAW_NEWLINE_JSON)
        self.assertTrue(extracted.startswith('def can_rearrange(s):'))
        self.assertIn('stack.append(char)', extracted)
        self.assertNotIn('function_code', extracted)
        compile(extracted, '<t>', 'exec')

    def test_observed_truncated_json_degrades_to_the_partial_source(self):
        extracted = _extract_code_block(OBSERVED_MBPP_TRUNCATED_JSON)
        self.assertTrue(extracted.startswith('def can_rearrange(s: str) -> bool:'))
        self.assertTrue(extracted.rstrip().endswith('return'))
        self.assertNotIn('"function"', extracted)

    def test_unfenced_json_object(self):
        extracted = _extract_code_block('{"code": "def f():\\n    return 7"}')
        self.assertEqual(extracted, 'def f():\n    return 7')

    def test_code_with_embedded_string_literals_survives_lenient_recovery(self):
        text = '```json\n{"function": "\ndef greet(name):\n    return "hello " + name\n"}\n```'
        extracted = _extract_code_block(text)
        self.assertTrue(extracted.startswith('def greet(name):'))
        self.assertIn('"hello "', extracted)

    def test_nested_json_object(self):
        text = '```json\n{"result": {"solution": "def f():\\n    return 3"}}\n```'
        self.assertEqual(_extract_code_block(text), 'def f():\n    return 3')

    def test_list_of_lines(self):
        text = '```json\n{"code": ["def f():", "    return 3"]}\n```'
        self.assertEqual(_extract_code_block(text), 'def f():\n    return 3')

    def test_non_code_sibling_keys_are_ignored(self):
        text = (
            '```json\n{"explanation": "This sorts the array in place.",\n'
            ' "function": "def f(xs):\\n    return sorted(xs)"}\n```'
        )
        self.assertEqual(_extract_code_block(text), 'def f(xs):\n    return sorted(xs)')

    def test_json_fence_holding_plain_python_is_left_to_the_legacy_path(self):
        # A ```json fence whose content is not a JSON object must not be
        # swallowed; behaviour there is unchanged from before the fix.
        text = '```json\ndef f():\n    return 1\n```'
        self.assertEqual(_extract_code_block(text), 'def f():\n    return 1')


class AmbiguousAndUnidentifiable(unittest.TestCase):
    """Extracting the wrong string is worse than extracting nothing."""

    def test_two_compiling_candidates_under_unknown_keys_yield_no_code(self):
        text = (
            '```json\n{"first_attempt": "def f():\\n    return 1",\n'
            ' "second_attempt": "def g():\\n    return 2"}\n```'
        )
        self.assertEqual(_extract_code_block(text), '')

    def test_two_candidates_disambiguate_by_a_conventional_key(self):
        text = (
            '```json\n{"scratchpad": "def draft():\\n    return 0",\n'
            ' "function": "def final():\\n    return 1"}\n```'
        )
        self.assertEqual(_extract_code_block(text), 'def final():\n    return 1')

    def test_compiling_candidate_beats_a_non_compiling_one(self):
        text = (
            '```json\n{"sketch": "def broken(:\\n    return",\n'
            ' "whatever": "def works():\\n    return 1"}\n```'
        )
        self.assertEqual(_extract_code_block(text), 'def works():\n    return 1')

    def test_json_object_without_any_python_yields_no_code(self):
        text = '```json\n{"answer": "B", "confidence": "0.9"}\n```'
        self.assertEqual(_extract_code_block(text), '')

    def test_empty_json_object_yields_no_code_but_does_not_crash(self):
        self.assertEqual(_extract_code_block('```json\n{}\n```'), '')

    def test_malformed_json_without_recoverable_fields_yields_no_code(self):
        self.assertEqual(_extract_code_block('```json\n{ garbage ][ \n```'), '')

    def test_truncated_json_before_any_value_yields_no_code(self):
        self.assertEqual(_extract_code_block('```json\n{\n  "function": '), '')

    def test_empty_and_non_string_input(self):
        self.assertEqual(_extract_code_block(''), '')
        self.assertEqual(_extract_code_block(None), '')


class EndToEndScoring(unittest.TestCase):
    """The extractor fix must move the actual benchmark score."""

    def test_json_wrapped_correct_solution_now_passes_execution(self):
        example = BenchmarkExample('humaneval', '1', 'code', 'q', {
            'tests': 'def check(candidate):\n    assert candidate(1, 2) == 3',
            'entry_point': 'add',
        })
        wrapped = '```json\n{"function": "def add(a, b):\\n    return a + b"}\n```'
        self.assertEqual(evaluate_code(wrapped, example, allow_exec=True), 1.0)

    def test_json_wrapped_wrong_solution_still_fails(self):
        example = BenchmarkExample('humaneval', '1', 'code', 'q', {
            'tests': 'def check(candidate):\n    assert candidate(1, 2) == 3',
            'entry_point': 'add',
        })
        wrapped = '```json\n{"function": "def add(a, b):\\n    return 0"}\n```'
        self.assertEqual(evaluate_code(wrapped, example, allow_exec=True), 0.0)

    def test_observed_humaneval_baseline_passes_the_real_tests(self):
        # HumanEval/109 gold tests, verbatim from the persisted result.json.
        example = BenchmarkExample('humaneval', 'HumanEval_109', 'code', 'q', {
            'tests': (
                'def check(candidate):\n'
                '    assert candidate([3, 4, 5, 1, 2])==True\n'
                '    assert candidate([3, 5, 10, 1, 2])==True\n'
                '    assert candidate([4, 3, 1, 2])==False\n'
                '    assert candidate([3, 5, 4, 1, 2])==False\n'
                '    assert candidate([])==True\n'
            ),
            'entry_point': 'move_one_ball',
        })
        self.assertEqual(evaluate_code(OBSERVED_HUMANEVAL_JSON, example, allow_exec=True), 1.0)


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