Claude
77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
"""Budget policy and accounting tests."""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import subprocess
|
|
|
|
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
|
sys.path.insert(0, os.path.join(ROOT, "pilot"))
|
|
|
|
from review.budget import Budget, BudgetState # noqa: E402
|
|
from review import opencode_runtime # noqa: E402
|
|
from review.opencode import parse_opencode_events # noqa: E402
|
|
|
|
|
|
def test_budget_reads_config_over_environment(monkeypatch):
|
|
monkeypatch.setenv("PRAGENT_MAX_REVIEW_STEPS", "3")
|
|
budget = Budget.from_config({"budget": {"max_steps": 7}})
|
|
assert budget.max_steps == 7
|
|
|
|
|
|
def test_budget_scales_for_broad_diff():
|
|
diff = "".join("+changed\n" for _ in range(850))
|
|
budget = Budget.for_review({}, diff)
|
|
assert budget.max_steps == 60
|
|
assert budget.max_total_tokens == 800_000
|
|
|
|
|
|
def test_explicit_budget_wins_over_diff_profile():
|
|
diff = "".join("+changed\n" for _ in range(2_100))
|
|
budget = Budget.for_review({"budget": {"max_steps": 9}}, diff)
|
|
assert budget.max_steps == 9
|
|
|
|
|
|
def test_budget_state_stops_at_token_limit():
|
|
state = BudgetState(Budget(max_steps=20, max_total_tokens=100))
|
|
assert state.record({"steps": 1, "total": 60, "output": 10}) == ""
|
|
assert state.record({"steps": 1, "total": 40, "output": 10}) == "max_total_tokens"
|
|
assert state.snapshot()["cap_hit"] is True
|
|
|
|
|
|
def test_budget_state_tracks_cost_cap():
|
|
state = BudgetState(Budget(max_equivalent_cost_usd=1.0))
|
|
assert state.record({"steps": 1, "total": 1}, 0.75) == ""
|
|
assert state.record({"steps": 1, "total": 1}, 0.25) == "max_equivalent_cost_usd"
|
|
|
|
|
|
def test_process_terminates_after_step_budget():
|
|
code = (
|
|
"import json,time; "
|
|
"print(json.dumps({'type':'step_finish','part':{'tokens':{"
|
|
"'input':1,'output':1,'total':2}}}), flush=True); "
|
|
"time.sleep(30)"
|
|
)
|
|
budget = Budget(max_steps=1, max_duration_seconds=10)
|
|
state = BudgetState(budget)
|
|
proc = opencode_runtime._run_process(
|
|
[sys.executable, "-u", "-c", code], cwd=".", env=os.environ.copy(),
|
|
timeout=10, parse_events=parse_opencode_events, budget=budget,
|
|
budget_state=state, model="glm-5.2:cloud", runner=subprocess.run,
|
|
)
|
|
assert state.snapshot()["cap_reason"] == "max_steps"
|
|
assert proc.stdout.count("step_finish") == 1
|
|
|
|
|
|
def test_process_terminates_silent_child_at_duration_budget():
|
|
code = "import time; time.sleep(30)"
|
|
budget = Budget(max_steps=20, max_duration_seconds=1)
|
|
state = BudgetState(budget)
|
|
proc = opencode_runtime._run_process(
|
|
[sys.executable, "-u", "-c", code], cwd=".", env=os.environ.copy(),
|
|
timeout=10, parse_events=parse_opencode_events, budget=budget,
|
|
budget_state=state, model="glm-5.2:cloud", runner=subprocess.run,
|
|
)
|
|
assert state.snapshot()["cap_reason"] == "max_duration_seconds"
|
|
assert proc.stdout == ""
|