"""Autoscaler 单测：深度计算、目标 worker 数、启停与缩容冷却。 只用临时 SQLite 与假 worker 命令（sleep 进程），不跑真实 worker、不联网、不碰生产库。假 worker 命令通过 worker_cmd 注入，scaler 的 worker_args 为空——真实 worker 命令构造 （强制 --loop/--worker-id/--worker-capacity）由构造测试覆盖。 """ import os import sqlite3 import sys import tempfile import time import unittest from pathlib import Path try: from cloud.autoscaler import Autoscaler, desired_workers, queue_depth from cloud.server import init_db except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.autoscaler import Autoscaler, desired_workers, queue_depth from cloud.server import init_db def _insert_task(db_path, task_id, status="queued", depends_on="[]"): now = time.time() conn = sqlite3.connect(db_path) try: conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " idempotency_key, depends_on, priority, created_by," " created_at, updated_at, last_heartbeat)" " VALUES(?, 'p1', 'x', ?, NULL, ?, 0, 'tester', ?, ?, NULL)", (task_id, status, depends_on, now, now), ) conn.commit() finally: conn.close() def _set_status(db_path, task_id, status): conn = sqlite3.connect(db_path) try: conn.execute("UPDATE tasks SET status = ? WHERE id = ?", (status, task_id)) conn.commit() finally: conn.close() class DesiredWorkersTestCase(unittest.TestCase): def test_depth_zero_keeps_min(self): self.assertEqual(desired_workers(0, 2, 1, 2), 1) def test_depth_zero_min_zero_scales_to_zero(self): self.assertEqual(desired_workers(0, 2, 0, 2), 0) def test_depth_divides_by_capacity_and_caps(self): self.assertEqual(desired_workers(1, 2, 1, 2), 1) self.assertEqual(desired_workers(3, 2, 1, 2), 2) self.assertEqual(desired_workers(10, 2, 1, 2), 2) self.assertEqual(desired_workers(5, 1, 1, 3), 3) class QueueDepthTestCase(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) def tearDown(self): self.tmp.cleanup() def test_only_ready_tasks_count(self): # A 就绪；B 依赖未完成的 A，不算；C 已完成，不算； # D 在跑，不算；E 失败，不算。 _insert_task(self.db_path, "A") _insert_task(self.db_path, "B", depends_on='["A"]') _insert_task(self.db_path, "C", status="completed") _insert_task(self.db_path, "D", status="running") _insert_task(self.db_path, "E", status="failed") self.assertEqual(queue_depth(self.db_path), 1) # A 完成后，B 解锁，深度仍是 1（不是 2）。 _set_status(self.db_path, "A", "completed") self.assertEqual(queue_depth(self.db_path), 1) def test_null_depends_on_counts_as_ready(self): # 回归：depends_on 为 NULL 时必须算无依赖，不能变成 # [None] 导致任务永远不算就绪（scaler 致盲）。 conn = sqlite3.connect(self.db_path) try: now = time.time() conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " idempotency_key, depends_on, priority, created_by," " created_at, updated_at, last_heartbeat)" " VALUES('N', 'p1', 'x', 'queued', NULL, NULL, 0," " 'tester', ?, ?, NULL)", (now, now)) conn.commit() finally: conn.close() self.assertEqual(queue_depth(self.db_path), 1) def test_dependency_cycle_counts_zero(self): _insert_task(self.db_path, "A", depends_on='["B"]') _insert_task(self.db_path, "B", depends_on='["A"]') self.assertEqual(queue_depth(self.db_path), 0) class AutoscalerReconcileTestCase(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) # 假 worker：睡 120 秒的 python 进程，scaler 用 SIGTERM # 停它（默认行为即退出，不用等 drain 超时）。 self.fake_cmd = [sys.executable, "-c", "import time; time.sleep(120)"] def tearDown(self): # 测试结束时残留的假 worker 必须清掉，否则泄漏进程。 if hasattr(self, "scaler"): for wid in list(self.scaler.workers): try: self.scaler._stop_worker(wid) except Exception: pass self.tmp.cleanup() def _scaler(self, **kwargs): kwargs.setdefault("min_workers", 1) kwargs.setdefault("max_workers", 2) kwargs.setdefault("per_worker_capacity", 2) kwargs.setdefault("scale_down_after", 1) kwargs.setdefault("drain_timeout", 5) self.scaler = Autoscaler( self.db_path, [], worker_cmd=self.fake_cmd, **kwargs) return self.scaler def test_scales_up_on_depth_and_down_after_cooldown(self): scaler = self._scaler() depth, desired, actual = scaler.reconcile() self.assertEqual((depth, desired, actual), (0, 1, 1)) for tid in ("A", "B", "C", "D"): _insert_task(self.db_path, tid) depth, desired, actual = scaler.reconcile() self.assertEqual((depth, desired, actual), (4, 2, 2)) # 任务全完成：scale_down_after=1，下一轮直接缩回 1。 for tid in ("A", "B", "C", "D"): _set_status(self.db_path, tid, "completed") depth, desired, actual = scaler.reconcile() self.assertEqual((depth, desired, actual), (0, 1, 1)) def test_scale_down_waits_for_cooldown_streak(self): scaler = self._scaler(scale_down_after=3) for tid in ("A", "B", "C", "D"): _insert_task(self.db_path, tid) scaler.reconcile() self.assertEqual(len(scaler.workers), 2) for tid in ("A", "B", "C", "D"): _set_status(self.db_path, tid, "completed") scaler.reconcile() self.assertEqual(len(scaler.workers), 2) scaler.reconcile() self.assertEqual(len(scaler.workers), 2) scaler.reconcile() self.assertEqual(len(scaler.workers), 1) def test_worker_command_forces_loop_id_capacity(self): scaler = Autoscaler(self.db_path, ["--interval", "10"]) cmd = scaler._worker_command("autoscale-1") self.assertIn("--loop", cmd) self.assertIn("autoscale-1", cmd) self.assertIn("--worker-capacity", cmd) self.assertEqual(cmd[cmd.index("--worker-capacity") + 1], "2") self.assertIn("--db", cmd) self.assertEqual(cmd[cmd.index("--db") + 1], self.db_path) def test_main_parses_args_and_starts_loop(self): # 回归：main() 里 setup_logging 曾错传 logger 名当 level， # 直接 ValueError，scaler 在主机上 crash-loop、一个 worker # 都没拉起来。本测试 mock 掉 run_loop 与 signal，只验证 # main() 能走到启动循环（即参数解析与日志初始化没炸）。 from unittest import mock import cloud.autoscaler as autoscaler_mod with mock.patch.object( autoscaler_mod.Autoscaler, "run_loop") as mock_loop, \ mock.patch("signal.signal"): autoscaler_mod.main([ "--db", self.db_path, "--interval", "30", "--", "--interval", "10"]) mock_loop.assert_called_once() if __name__ == "__main__": unittest.main()