"""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() def _insert_running_task(db_path, task_id, parts=None, heartbeat_offset=0): """插入一个 running 任务；parts 为 coordinator_parts 列的 原始值（None=普通任务，JSON 文本=协调员分工）。""" _insert_task(db_path, task_id, status="running") now = time.time() conn = sqlite3.connect(db_path) try: conn.execute( "UPDATE tasks SET coordinator_parts = ?," " last_heartbeat = ? WHERE id = ?", (parts, now + heartbeat_offset, task_id)) conn.commit() finally: conn.close() def _set_worker_running(db_path, worker_id, running): """在 workers 表里给 worker 写一行新鲜心跳 + running 数。 用 cloud.workers 的注册函数建表，避免测试自己拼表结构。""" try: from cloud.workers import register_worker except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.workers import register_worker register_worker(db_path, worker_id, executor="docker", capacity=2, running=running) 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() class ScaleDownGuardTestCase(unittest.TestCase): """缩容护栏单测：两层 veto 与 streak 冻结。 全部用临时 SQLite + 假 worker（sleep 进程），不跑真实 worker。注意：假 worker 没有 workers 表心跳，注册行 全部由 _set_worker_running 显式写；没写注册行的 worker 在护栏里按空闲处理（见 autoscaler docstring）。 """ def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") init_db(self.db_path) self.fake_cmd = [sys.executable, "-c", "import time; time.sleep(120)"] def tearDown(self): 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("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 _scale_to(self, scaler, count): for i in range(count * 2): _insert_task(self.db_path, "Q%d" % i) scaler.reconcile() self.assertEqual(len(scaler.workers), count) for i in range(count * 2): _set_status(self.db_path, "Q%d" % i, "completed") def test_live_coordinator_blocks_scale_down(self): scaler = self._scaler(max_workers=2) self._scale_to(scaler, 2) _insert_running_task( self.db_path, "COORD", parts='[{"name": "A", "prompt": "x"}]') scaler.reconcile() self.assertEqual(len(scaler.workers), 2) self.assertEqual(scaler.last_action, "scale_down_blocked_by_coordinator") def test_stale_coordinator_does_not_block(self): scaler = self._scaler(max_workers=2) self._scale_to(scaler, 2) _insert_running_task( self.db_path, "COORD", parts='[{"name": "A", "prompt": "x"}]', heartbeat_offset=-300) scaler.reconcile() self.assertEqual(len(scaler.workers), 1) def test_non_coordinator_running_does_not_trigger_coordinator_veto(self): scaler = self._scaler(max_workers=2) self._scale_to(scaler, 2) _insert_running_task(self.db_path, "PLAIN", parts=None) _insert_running_task(self.db_path, "EMPTY", parts="[]") scaler.reconcile() self.assertEqual(len(scaler.workers), 1) self.assertEqual(scaler.last_action, "scaled_down_to_1") def test_invalid_coordinator_parts_not_a_coordinator(self): # 与 worker 侧一致：parse_parts 把坏值解析为 []， # worker 不会把它当协调员跑等齐循环，护栏也不拦。 scaler = self._scaler(max_workers=2) self._scale_to(scaler, 2) _insert_running_task(self.db_path, "BAD", parts="not-json") scaler.reconcile() self.assertEqual(len(scaler.workers), 1) def test_busy_victim_skipped_idle_victim_stopped(self): scaler = self._scaler(max_workers=3) self._scale_to(scaler, 3) # DB 里有一个活 running 任务（幽灵交叉校验的另一半）。 _insert_running_task(self.db_path, "LIVE", parts=None) _set_worker_running(self.db_path, "autoscale-1", 0) _set_worker_running(self.db_path, "autoscale-2", 0) _set_worker_running(self.db_path, "autoscale-3", 1) scaler.reconcile() # LIFO 受害者是 [3, 2]：3 忙跳过，2 闲停掉，实际 3->2。 self.assertEqual(len(scaler.workers), 2) self.assertIn("autoscale-3", scaler.workers) self.assertNotIn("autoscale-2", scaler.workers) self.assertEqual(scaler.last_action, "scaled_down_to_2") def test_ghost_registry_does_not_block(self): scaler = self._scaler(max_workers=2) self._scale_to(scaler, 2) # 注册表报 busy，但 DB 里一个活 running 任务都没有： # 幽灵占位，不拦。 _set_worker_running(self.db_path, "autoscale-2", 2) scaler.reconcile() self.assertEqual(len(scaler.workers), 1) def test_coordinator_veto_freezes_low_streak(self): scaler = self._scaler(max_workers=2, scale_down_after=3) self._scale_to(scaler, 2) _insert_running_task( self.db_path, "COORD", parts='[{"name": "A", "prompt": "x"}]') for _ in range(3): scaler.reconcile() self.assertEqual(len(scaler.workers), 2) # veto 解除后不需要再等 3 轮：streak 已冻结在阈值上， # 下一轮低深度立即缩。 _set_status(self.db_path, "COORD", "completed") scaler.reconcile() self.assertEqual(len(scaler.workers), 1) if __name__ == "__main__": unittest.main()