"""自动化执行器入口 run_due_automation:抢单 + 配额门单测。""" import os import sqlite3 import tempfile import time import unittest from cloud.auth import audit_on_conn from cloud.automations import ( create_automation, ensure_automations_tables, list_automation_runs, run_due_automation, ) from cloud.ledger import GENESIS_HASH, compute_entry_hash 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 AuditPairPgIdentityTestCase(unittest.TestCase): """模拟 PG:id 是 BIGINT identity,显式插 TEXT uuid 会报错。 假连接只接受不带 id 的 INSERT;旧 _audit_row(带 TEXT id) 在此假连接上应失败。接替它的 audit_on_conn 须同时做到: audit_log 行不带 id、链条目与该行同 created_at、空链时 prev 为 GENESIS、entry_hash 按规范载荷重算一致;且本函数 不 commit(事务属于调用方,见 run_due_automation 文档)。 假连接类名须为 PgConnection:ledger.ensure_schema 靠类名 判定 PG 连接以跳过 SQLite 触发器,与真实 PG 适配层同款。 """ class _FakeCursor: def __init__(self, row=None): self._row = row self.lastrowid = 1 def fetchone(self): return self._row class PgConnection: def __init__(self): self.audit_rows = [] self.ledger_rows = [] self.committed = False def execute(self, sql, params=()): upper = sql.strip().upper() if upper.startswith("INSERT INTO AUDIT_LOG(ID"): raise ValueError("column \"id\" is of type bigint") if upper.startswith("INSERT INTO AUDIT_LOG"): self.audit_rows.append(params) return AuditPairPgIdentityTestCase._FakeCursor() if upper.startswith("SELECT ENTRY_HASH FROM AUDIT_LEDGER"): return AuditPairPgIdentityTestCase._FakeCursor(None) if upper.startswith("INSERT INTO AUDIT_LEDGER"): self.ledger_rows.append(params) return AuditPairPgIdentityTestCase._FakeCursor() if upper.startswith("CREATE TABLE"): return AuditPairPgIdentityTestCase._FakeCursor() raise ValueError("unexpected sql in fake PG conn: %r" % sql) def commit(self): self.committed = True def rollback(self): pass def test_pair_without_id_succeeds_on_pg_like_conn(self): conn = self.PgConnection() audit_on_conn(conn, "system", "p1", "automation_run_created", "auto1 -> task1", created_at=123.0) # 调用方拥有事务:本函数自己不 commit。 self.assertFalse(conn.committed) self.assertEqual(len(conn.audit_rows), 1) # 参数里没有 uuid:第一位是 system,第二位是 project_id。 self.assertEqual(conn.audit_rows[0][0], "system") self.assertEqual(conn.audit_rows[0][1], "p1") self.assertEqual(conn.audit_rows[0][4], 123.0) self.assertEqual(len(conn.ledger_rows), 1) ledger = conn.ledger_rows[0] # audit_ledger 列序:user_id, project_id, task_id, action, # detail, created_at, prev_hash, entry_hash。 self.assertEqual(ledger[0], "system") self.assertEqual(ledger[1], "p1") self.assertIsNone(ledger[2]) self.assertEqual(ledger[3], "automation_run_created") self.assertEqual(ledger[4], "auto1 -> task1") self.assertEqual(ledger[5], 123.0) self.assertEqual(ledger[6], GENESIS_HASH) self.assertEqual( ledger[7], compute_entry_hash(GENESIS_HASH, "system", "p1", None, "automation_run_created", "auto1 -> task1", 123.0)) 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()