rein-openweights/tests/test_loop.py

import subprocess
from unittest.mock import MagicMock

from rein_openweights.budget import BudgetTracker
from rein_openweights.loop import run_loop


def _init_repo(tmp_path):
    subprocess.run(["git", "init", "-q"], cwd=tmp_path, check=True)
    subprocess.run(["git", "config", "user.email", "t@example.com"], cwd=tmp_path, check=True)
    subprocess.run(["git", "config", "user.name", "t"], cwd=tmp_path, check=True)
    subprocess.run(["git", "commit", "-q", "--allow-empty", "-m", "init"], cwd=tmp_path, check=True)
    return tmp_path


def _response(message, total_tokens=100):
    return {
        "choices": [{"message": message}],
        "usage": {"total_tokens": total_tokens},
    }


def test_run_loop_executes_tool_call_then_finishes(tmp_path):
    repo = _init_repo(tmp_path)
    client = MagicMock()
    client.chat.side_effect = [
        _response(
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "call_1",
                        "function": {
                            "name": "write_file",
                            "arguments": '{"path": "NOTES.md", "content": "hi"}',
                        },
                    }
                ],
            }
        ),
        _response(
            {
                "role": "assistant",
                "tool_calls": [
                    {
                        "id": "call_2",
                        "function": {
                            "name": "git_add_commit",
                            "arguments": '{"message": "add notes"}',
                        },
                    }
                ],
            }
        ),
        _response({"role": "assistant", "content": "All set. DONE"}),
    ]

    result = run_loop(client, repo, "t", "d", budget=BudgetTracker(total=10_000))

    assert result.finished is True
    assert result.turns == 3
    assert (repo / "NOTES.md").read_text() == "hi"
    log = subprocess.run(
        ["git", "-C", str(repo), "log", "--oneline"], capture_output=True, text=True
    ).stdout
    assert "add notes" in log


def test_run_loop_stops_when_budget_exceeded(tmp_path):
    repo = _init_repo(tmp_path)
    client = MagicMock()
    client.chat.return_value = _response({"role": "assistant", "content": "still going"}, total_tokens=100)

    result = run_loop(client, repo, "t", "d", budget=BudgetTracker(total=50))

    assert result.error is not None
    assert "budget" in result.error.lower()


def test_run_loop_stops_at_max_turns(tmp_path):
    repo = _init_repo(tmp_path)
    client = MagicMock()
    client.chat.return_value = _response(
        {
            "role": "assistant",
            "tool_calls": [
                {"id": "c", "function": {"name": "git_status", "arguments": "{}"}}
            ],
        }
    )

    result = run_loop(client, repo, "t", "d", max_turns=2)

    assert result.turns == 2
    assert result.error == "max_turns exceeded"


def test_run_loop_reports_tool_error_without_crashing(tmp_path):
    repo = _init_repo(tmp_path)
    client = MagicMock()
    client.chat.side_effect = [
        _response(
            {
                "role": "assistant",
                "tool_calls": [
                    {"id": "c", "function": {"name": "read_file", "arguments": '{"path": "missing.txt"}'}}
                ],
            }
        ),
        _response({"role": "assistant", "content": "done anyway DONE"}),
    ]

    result = run_loop(client, repo, "t", "d")

    assert result.finished is True
    tool_messages = [m for m in result.transcript if m.get("role") == "tool"]
    assert "error" in tool_messages[0]["content"]