test_budget_guardrail.py
4.18 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
"""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 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")
if __name__ == "__main__":
unittest.main()