"""自动化执行器入口 run_due_automation:抢单 + 配额门单测。""" import os import sqlite3 import tempfile import time import unittest from cloud.automations import ( _audit_row, create_automation, ensure_automations_tables, list_automation_runs, run_due_automation, ) from cloud.server import init_db def _tasks(db_path): conn = sqlite3.connect(db_path) try: return conn.execute("SELECT id, prompt, status FROM tasks").fetchall() finally: conn.close() def _audit_count(db_path, action): conn = sqlite3.connect(db_path) try: return conn.execute( "SELECT COUNT(*) FROM audit_log WHERE action = ?", (action,)).fetchone()[0] finally: conn.close() def _set_quota(db_path, user_id, micros): conn = sqlite3.connect(db_path) try: conn.execute( "CREATE TABLE IF NOT EXISTS user_quotas(" "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, " " updated_at DOUBLE PRECISION)") conn.execute( "INSERT OR REPLACE INTO user_quotas(user_id," " quota_usd_micros, updated_at) VALUES(?, ?, ?)", (user_id, micros, time.time())) conn.commit() finally: conn.close() class AuditRowPgIdentityTestCase(unittest.TestCase): """模拟 PG:id 是 BIGINT identity,显式插 TEXT uuid 会报错。 假连接只接受不带 id 的 INSERT;旧实现(带 TEXT id)在此 假连接上应失败、新实现应成功且审计行落盘。 """ class _PgLikeConn: def __init__(self): self.rows = [] self.committed = False def execute(self, sql, params=()): if sql.strip().upper().startswith("INSERT INTO AUDIT_LOG(ID"): raise ValueError("column \"id\" is of type bigint") if sql.strip().upper().startswith("INSERT INTO AUDIT_LOG"): self.rows.append(params) return None raise ValueError("unexpected sql in fake PG conn") def commit(self): self.committed = True def rollback(self): pass def test_insert_without_id_succeeds_on_pg_like_conn(self): conn = self._PgLikeConn() _audit_row(conn, "p1", "automation_run_created", "auto1 -> task1", 123.0) self.assertTrue(conn.committed) self.assertEqual(len(conn.rows), 1) # 参数里没有 uuid:第一位是 system,第二位是 project_id。 self.assertEqual(conn.rows[0][0], "system") self.assertEqual(conn.rows[0][1], "p1") class RunDueAutomationTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") init_db(self.db_path) self.conn = sqlite3.connect(self.db_path) self.conn.row_factory = sqlite3.Row ensure_automations_tables(self.conn) def tearDown(self): self.conn.close() self.tmp.cleanup() def _auto(self, owner="owner1"): return create_automation(self.conn, "proj", "daily", "do thing", owner, interval_seconds=60) def test_run_creates_task_within_quota(self): auto = self._auto() run = run_due_automation(self.conn, auto["id"]) self.assertIsNotNone(run) self.assertEqual(run["status"], "created") self.assertIsNotNone(run["task_id"]) tasks = _tasks(self.db_path) self.assertEqual(len(tasks), 1) self.assertEqual(tasks[0][1], "do thing") self.assertEqual(tasks[0][2], "queued") self.assertEqual( _audit_count(self.db_path, "automation_run_created"), 1) # 同一间隔内第二次跑:已抢过,返回 None,不重复建任务。 self.assertIsNone(run_due_automation(self.conn, auto["id"])) self.assertEqual(len(_tasks(self.db_path)), 1) def test_run_skips_when_over_quota(self): auto = self._auto() _set_quota(self.db_path, "owner1", 0) run = run_due_automation(self.conn, auto["id"]) self.assertIsNotNone(run) self.assertEqual(run["status"], "skipped_quota") self.assertIsNone(run["task_id"]) self.assertEqual(_tasks(self.db_path), []) self.assertEqual( _audit_count(self.db_path, "automation_run_quota_skipped"), 1) # 跳过仍记了产出(不是失败),同间隔内不再重跑。 runs = list_automation_runs(self.conn, auto["id"]) self.assertEqual(len(runs), 1) self.assertEqual(runs[0]["status"], "skipped_quota") def test_run_not_due_returns_none(self): auto = self._auto() self.conn.execute( "UPDATE automations SET last_run_at = ? WHERE id = ?", (time.time(), auto["id"])) self.conn.commit() self.assertIsNone(run_due_automation(self.conn, auto["id"])) self.assertEqual(_tasks(self.db_path), []) self.assertEqual( list_automation_runs(self.conn, auto["id"]), []) if __name__ == "__main__": unittest.main()