codex-app-server/cloud/tests/test_automation_run_due.py

"""自动化执行器入口 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()