tests/test_recursion_extract.py

"""Tests for rewards.recursion_extract — REW-04 v2."""
from rewards.recursion_extract import count_llm_calls


def test_empty_completion():
    assert count_llm_calls("") == 0


def test_no_code_block():
    assert count_llm_calls("The answer is <answer>silver</answer>") == 0


def test_single_call():
    c = "```python\nresult = llm('find', ctx[:5000])\n```\n<answer>silver</answer>"
    assert count_llm_calls(c) == 1


def test_two_calls_in_one_block():
    c = "```python\na = llm('q1', ctx[:1000])\nb = llm('q2', ctx[1000:])\n```"
    assert count_llm_calls(c) == 2


def test_calls_across_two_blocks():
    c = "```python\nx=llm('q', ctx)\n```\nthen\n```python\ny=llm('q2', ctx)\n```"
    assert count_llm_calls(c) == 2


def test_call_in_comment_not_counted():
    c = "```python\n# would call llm(stuff) but skipping\nprint('done')\n```"
    assert count_llm_calls(c) == 0


def test_call_in_string_literal_not_counted():
    c = '```python\nnote = "earlier code did llm(...)"\nprint(note)\n```'
    assert count_llm_calls(c) == 0


def test_call_outside_code_block_not_counted():
    c = "Maybe I should call llm(question, chunk) but I won't actually."
    assert count_llm_calls(c) == 0


def test_call_in_loop_counts_literal_occurrence():
    c = "```python\nfor chunk in chunks:\n    r = llm('find', chunk)\n```"
    assert count_llm_calls(c) == 1


def test_invalid_python_falls_back_to_regex():
    c = "```python\nthis is not valid python !!!\nresult = llm('q', ctx)\n```"
    assert count_llm_calls(c) >= 1  # fallback regex finds it


def test_bare_python_fence():
    c = "```\nans = llm('q', ctx)\n```"  # no language tag
    assert count_llm_calls(c) == 1