test_budget_guardrail.py
6.74 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
"""U-019 tests: c_budget_check + update_budget_limit_and_resume. All real
(sqlite is fast/free/local, no mocking needed).
Run: python3 -m unittest src.test_budget_guardrail -v
"""
import unittest
from src import budget_guardrail, schema
class BudgetGuardrailTestBase(unittest.TestCase):
def setUp(self):
self.conn = schema.connect(":memory:")
schema.create_schema(self.conn)
self.conn.execute(
"INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
"VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
)
self.conn.execute(
"INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) "
"VALUES ('task-1', 'proj-1', 'goal', 'goal', 'coding', 'coder')"
)
self.conn.execute(
"INSERT INTO budget_guardrail_config (config_id, scope, limit_type, limit_value, current_usage, "
"updated_by) VALUES ('default', 'global', 'token_count', 100, 0, 'test')"
)
self.conn.commit()
def tearDown(self):
self.conn.close()
class ListBudgetConfigsTest(BudgetGuardrailTestBase):
def test_returns_the_configured_row(self):
result = budget_guardrail.list_budget_configs(self.conn)
self.assertEqual(len(result), 1)
self.assertEqual(
result[0],
{
"config_id": "default", "scope": "global", "limit_type": "token_count",
"limit_value": 100, "current_usage": 0, "updated_at": result[0]["updated_at"],
},
)
def test_no_config_returns_empty_list(self):
self.conn.execute("DELETE FROM budget_guardrail_config")
self.conn.commit()
self.assertEqual(budget_guardrail.list_budget_configs(self.conn), [])
class CheckBudgetTest(BudgetGuardrailTestBase):
def test_under_limit_allowed_and_accumulates(self):
result = budget_guardrail.check_budget(self.conn, "task-1", "default", 30)
self.assertEqual(result, {"allowed": True, "remaining_quota": 70})
usage = self.conn.execute("SELECT current_usage FROM budget_guardrail_config WHERE config_id = 'default'").fetchone()[0]
self.assertEqual(usage, 30)
def test_reaching_limit_blocks_task_and_raises(self):
# REQ-10 happy: small limit, run past it -> blocked_pending_human,
# blocked_reason='budget_exceeded'
with self.assertRaises(budget_guardrail.BudgetExceededError):
budget_guardrail.check_budget(self.conn, "task-1", "default", 150)
row = self.conn.execute("SELECT status, blocked_reason FROM task WHERE task_id = 'task-1'").fetchone()
self.assertEqual(row, ("blocked_pending_human", "budget_exceeded"))
def test_exceeded_check_does_not_accumulate_usage(self):
with self.assertRaises(budget_guardrail.BudgetExceededError):
budget_guardrail.check_budget(self.conn, "task-1", "default", 150)
usage = self.conn.execute("SELECT current_usage FROM budget_guardrail_config WHERE config_id = 'default'").fetchone()[0]
self.assertEqual(usage, 0)
def test_blocked_reason_distinct_from_usage_limit_wait_timeout(self):
# REQ-10 failure: the two guardrails must never share a blocked_reason
with self.assertRaises(budget_guardrail.BudgetExceededError):
budget_guardrail.check_budget(self.conn, "task-1", "default", 150)
blocked_reason = self.conn.execute("SELECT blocked_reason FROM task WHERE task_id = 'task-1'").fetchone()[0]
self.assertNotEqual(blocked_reason, "usage_limit_wait_timeout")
self.assertEqual(blocked_reason, "budget_exceeded")
class ResumeAfterLimitRaiseTest(BudgetGuardrailTestBase):
def test_raising_limit_resumes_budget_blocked_tasks(self):
with self.assertRaises(budget_guardrail.BudgetExceededError):
budget_guardrail.check_budget(self.conn, "task-1", "default", 150)
result = budget_guardrail.update_budget_limit_and_resume(self.conn, "default", 1000, "admin")
self.assertEqual(result["resumed_task_ids"], ["task-1"])
row = self.conn.execute("SELECT status, blocked_reason FROM task WHERE task_id = 'task-1'").fetchone()
self.assertEqual(row, ("coding", None)) # restored to pre-blocked status ('coding')
def test_does_not_resume_tasks_blocked_for_other_reasons(self):
self.conn.execute(
"UPDATE task SET status = 'blocked_pending_human', blocked_reason = 'loop_limit_reached' "
"WHERE task_id = 'task-1'"
)
self.conn.commit()
result = budget_guardrail.update_budget_limit_and_resume(self.conn, "default", 1000, "admin")
self.assertEqual(result["resumed_task_ids"], [])
status = self.conn.execute("SELECT status FROM task WHERE task_id = 'task-1'").fetchone()[0]
self.assertEqual(status, "blocked_pending_human")
class UpsertBudgetConfigTest(BudgetGuardrailTestBase):
def test_updates_existing_config_preserving_current_usage(self):
self.conn.execute("UPDATE budget_guardrail_config SET current_usage = 42 WHERE config_id = 'default'")
self.conn.commit()
result = budget_guardrail.upsert_budget_config(self.conn, "global", "token_count", 500, "admin")
self.assertEqual(result["config"]["limit_value"], 500)
self.assertEqual(result["config"]["current_usage"], 42)
def test_creates_config_when_none_exists(self):
self.conn.execute("DELETE FROM budget_guardrail_config")
self.conn.commit()
result = budget_guardrail.upsert_budget_config(self.conn, "per_task", "call_count", 10, "admin")
self.assertEqual(result["config"]["current_usage"], 0)
self.assertEqual(result["config"]["scope"], "per_task")
def test_resumes_blocked_tasks_same_as_update_budget_limit(self):
with self.assertRaises(budget_guardrail.BudgetExceededError):
budget_guardrail.check_budget(self.conn, "task-1", "default", 150)
result = budget_guardrail.upsert_budget_config(self.conn, "global", "token_count", 1000, "admin")
self.assertEqual(result["resumed_task_ids"], ["task-1"])
def test_invalid_scope_rejected(self):
with self.assertRaises(budget_guardrail.InvalidParamError):
budget_guardrail.upsert_budget_config(self.conn, "not_a_scope", "token_count", 100, "admin")
def test_invalid_limit_type_rejected(self):
with self.assertRaises(budget_guardrail.InvalidParamError):
budget_guardrail.upsert_budget_config(self.conn, "global", "not_a_type", 100, "admin")
def test_zero_or_negative_limit_value_rejected(self):
with self.assertRaises(budget_guardrail.InvalidParamError):
budget_guardrail.upsert_budget_config(self.conn, "global", "token_count", 0, "admin")
if __name__ == "__main__":
unittest.main()