| from local_train.think import thought_text, wrap_completion | |
| def test_wrap_converts_thought_prefix_and_keeps_one_bash_block(): | |
| gold = "THOUGHT: Look at retention.go\n\n```bash\ncat /testbed/retention.go\n```\n```bash\nls\n```" | |
| out = wrap_completion(gold) | |
| assert out is not None | |
| assert out.startswith("<think>\nLook at retention.go") | |
| assert "</think>" in out | |
| assert out.count("```bash") == 1 | |
| assert "cat /testbed/retention.go" in out | |
| assert "ls\n```" not in out | |
| def test_wrap_override_rewrites_submit_command(): | |
| gold = "THOUGHT: done\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```" | |
| out = wrap_completion(gold, "echo ALBEDO_TASK_DONE_SUBMIT_NOW") | |
| assert out is not None | |
| assert "echo ALBEDO_TASK_DONE_SUBMIT_NOW" in out | |
| assert "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" not in out.split("</think>", 1)[-1] | |
| def test_wrap_rejects_missing_bash(): | |
| assert wrap_completion("THOUGHT: just thinking") is None | |
| def test_existing_think_block_is_reused(): | |
| gold = "<think>plan</think>\n\n```bash\nsed -i 's/a/b/' foo.py\n```" | |
| assert thought_text(gold) == "plan" | |
| out = wrap_completion(gold) | |
| assert out is not None | |
| assert "<think>\nplan\n</think>" in out | |