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