"""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()