code/tests/test_local_train_reward.py

from local_train.reward import reward_completions, score_texts


def test_fatal_loop_and_missing_bash_are_zero():
    looped = ["```bash\nls\n```"] * 4
    assert score_texts(looped, submit_command="echo X").reward == 0.0
    assert score_texts(["no command here"], submit_command="echo X").fatal


def test_exact_submit_after_edit_scores_high():
    texts = [
        "<think>edit</think>\n\n```bash\nsed -i 's/a/b/' src/foo.py\n```",
        "<think>submit</think>\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW\n```",
    ]
    got = score_texts(
        texts,
        submit_command="echo ALBEDO_TASK_DONE_SUBMIT_NOW",
        gold_paths=["src/foo.py"],
    )
    assert got.edited and got.submitted and not got.submit_before_edit
    assert got.reward >= 0.7


def test_marker_submit_without_exact_tail_counts():
    texts = [
        "<think>edit</think>\n\n```bash\nsed -i 's/a/b/' src/foo.py\n```",
        "<think>s</think>\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW && cat leftover.txt\n```",
    ]
    got = score_texts(
        texts,
        submit_command="echo ALBEDO_TASK_DONE_SUBMIT_NOW && git add -A && git diff --cached",
        gold_paths=["src/foo.py"],
    )
    assert got.submitted and not got.submit_before_edit


def test_submit_before_edit_is_penalized():
    texts = ["<think>x</think>\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW\n```"]
    got = score_texts(texts, submit_command="echo ALBEDO_TASK_DONE_SUBMIT_NOW")
    assert got.submitted and got.submit_before_edit
    assert got.reward < 0.4


def test_dummy_edit_is_penalized():
    texts = ["<think>x</think>\n\n```bash\ncat <<'EOF' > newfile.py\nprint(1)\nEOF\n```"]
    got = score_texts(texts)
    assert got.dummy
    assert got.reward < 0.2


def test_grpo_callback_aligns_extra_columns():
    completions = [
        "<think>x</think>\n\n```bash\necho FINALIZE_AND_SUBMIT_TASK_OUTPUT\n```",
        "<think>x</think>\n\n```bash\nsed -i 's/a/b/' lib/bar.py\n```",
    ]
    rewards = reward_completions(
        completions,
        submit_command=["echo FINALIZE_AND_SUBMIT_TASK_OUTPUT", "echo FINALIZE_AND_SUBMIT_TASK_OUTPUT"],
        gold_paths=[["lib/bar.py"], ["lib/bar.py"]],
    )
    assert rewards[1] > rewards[0]