"""生产级缺口单测：worker 常驻/真取消、tasks 迁移、metrics、备份、日志。 只用标准库，不跑真实模型回合、不联网、不登录任何服务； process_factory 返回的长跑子进程只是 `sleep`，取消后必须消失。 """ import http.client import io import json import logging import os import sqlite3 import subprocess import sys import tempfile import threading import time import unittest from unittest import mock from pathlib import Path from types import SimpleNamespace try: from cloud.auth import AuthStore from cloud.backup import backup_db from cloud.cost import CostStore from cloud.db import PostgresDriverMissingError from cloud.logging_config import JsonFormatter, log_event, setup_logging from cloud.server import create_server, init_db from cloud.worker import Worker except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.auth import AuthStore from cloud.backup import backup_db from cloud.cost import CostStore from cloud.db import PostgresDriverMissingError from cloud.logging_config import JsonFormatter, log_event, setup_logging from cloud.server import create_server, init_db from cloud.worker import Worker def _insert_task(db_path, task_id, prompt="do work", status="queued", project_id="p1", started=None, finished=None): 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(?, ?, ?, ?, NULL, '[]', 0, 'tester', ?, ?, NULL)", (task_id, project_id, prompt, status, now, now), ) if started is not None: conn.execute( "UPDATE tasks SET started_at = ?, finished_at = ?" " WHERE id = ?", (started, finished, task_id)) conn.commit() finally: conn.close() def _row(db_path, task_id): conn = sqlite3.connect(db_path) conn.row_factory = sqlite3.Row try: row = conn.execute( "SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone() return dict(row) if row else None finally: conn.close() def _completed_runner(cmd): return SimpleNamespace( returncode=0, stdout=json.dumps({"task": "x", "status": "completed", "text": "done"}), stderr="") class MigrationTestCase(unittest.TestCase): def test_old_db_migrates_without_data_loss(self): tmp = tempfile.TemporaryDirectory() db_path = os.path.join(tmp.name, "old.db") conn = sqlite3.connect(db_path) try: conn.execute( "CREATE TABLE projects(id TEXT PRIMARY KEY, name TEXT," " owner_id TEXT, created_at REAL)") conn.execute( "CREATE TABLE tasks(id TEXT PRIMARY KEY, project_id TEXT," " prompt TEXT, status TEXT, idempotency_key TEXT," " depends_on TEXT, priority INTEGER, created_by TEXT," " created_at REAL, updated_at REAL, last_heartbeat REAL," " UNIQUE(project_id, idempotency_key))") conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " depends_on, priority, created_by, created_at," " updated_at, last_heartbeat)" " VALUES('old1', 'p1', 'old prompt', 'completed', '[]'," " 0, 'tester', 1, 1, NULL)") conn.commit() finally: conn.close() init_db(db_path) cols = {r[1] for r in sqlite3.connect(db_path).execute( "PRAGMA table_info(tasks)").fetchall()} for col in ("result_text", "error_text", "actual_model", "started_at", "finished_at", "child_pid"): self.assertIn(col, cols) row = _row(db_path, "old1") self.assertEqual(row["prompt"], "old prompt") self.assertEqual(row["status"], "completed") # 再跑一次 init_db 是 no-op，不炸、不丢数据。 init_db(db_path) self.assertEqual(_row(db_path, "old1")["prompt"], "old prompt") tmp.cleanup() class WorkerResultColumnsTestCase(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_completed_task_writes_result_columns(self): _insert_task(self.db_path, "t1") worker = Worker(self.db_path, "/fake/host.py", runner=_completed_runner, env={"XCM_API_KEY": "test-placeholder"}) result = worker.run_one() self.assertEqual(result["status"], "completed") row = _row(self.db_path, "t1") self.assertIn("done", row["result_text"]) self.assertIsNone(row["error_text"]) self.assertEqual(row["actual_model"], worker.model) self.assertIsNotNone(row["started_at"]) self.assertIsNotNone(row["finished_at"]) self.assertGreaterEqual(row["finished_at"], row["started_at"]) def test_missing_key_writes_error_text_and_calls_runner_zero_times(self): calls = [] def counting(cmd): calls.append(cmd) return _completed_runner(cmd) _insert_task(self.db_path, "t2") worker = Worker(self.db_path, "/fake/host.py", runner=counting, env={}) result = worker.run_one() self.assertEqual(result["status"], "failed") self.assertEqual(calls, []) row = _row(self.db_path, "t2") self.assertIn("Missing environment variable", row["error_text"]) self.assertIsNotNone(row["started_at"]) class WorkerTrueCancelTestCase(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 _sleep_factory(self, task): return subprocess.Popen( [sys.executable, "-c", "import time; time.sleep(30)"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True) def test_cancel_kills_process_group_and_keeps_interrupted(self): _insert_task(self.db_path, "long1") worker = Worker(self.db_path, "/fake/host.py", process_factory=self._sleep_factory, env={"XCM_API_KEY": "test-placeholder"}) results = [] thread = threading.Thread( target=lambda: results.append(worker.run_one()), daemon=True) thread.start() pgid = None deadline = time.time() + 5 while time.time() < deadline: row = _row(self.db_path, "long1") if row and row["child_pid"]: pgid = row["child_pid"] break time.sleep(0.05) self.assertIsNotNone(pgid, "子进程 pid 未写进任务记录") # 进程组当前存在（killpg 0 不抛）。 os.killpg(os.getpgid(pgid), 0) self.assertTrue(worker.cancel("long1")) thread.join(timeout=10) self.assertFalse(thread.is_alive()) self.assertEqual(results[0]["status"], "interrupted") self.assertEqual(_row(self.db_path, "long1")["status"], "interrupted") # 进程组确实消失。 with self.assertRaises((ProcessLookupError, PermissionError, OSError)): os.killpg(os.getpgid(pgid), 0) def test_loop_stops_and_marks_running_interrupted(self): _insert_task(self.db_path, "loop1") worker = Worker(self.db_path, "/fake/host.py", process_factory=self._sleep_factory, env={"XCM_API_KEY": "test-placeholder"}) stop = threading.Event() results = [] thread = threading.Thread( target=lambda: results.extend( worker.run_loop(interval=0.2, stop_event=stop)), daemon=True) thread.start() time.sleep(1.0) stop.set() # 循环在任务进行中收到停止信号：_loop_stop 由 run_loop 的 # finally 外的轮询感知——这里直接置位模拟 SIGTERM 处理。 worker._loop_stop = True thread.join(timeout=10) self.assertFalse(thread.is_alive()) self.assertEqual(_row(self.db_path, "loop1")["status"], "interrupted") class MetricsAndBackupTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") self.auth = AuthStore(self.db_path) self.owner_id, self.owner_key = self.auth.create_user("owner") self.server = create_server(self.db_path, "127.0.0.1", 0) self.port = self.server.server_address[1] self.thread = threading.Thread( target=self.server.serve_forever, daemon=True) self.thread.start() def tearDown(self): self.server.shutdown() self.server.server_close() self.thread.join(timeout=5) self.tmp.cleanup() def _req(self, method, path, key=None, body=None): conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10) hdrs = {} payload = None if body is not None: payload = json.dumps(body).encode() hdrs["Content-Type"] = "application/json" if key: hdrs["X-API-Key"] = key try: conn.request(method, path, body=payload, headers=hdrs) resp = conn.getresponse() data = resp.read() finally: conn.close() return resp.status, json.loads(data) if data else None def test_metrics_counts_match_db(self): status, proj = self._req("POST", "/projects", self.owner_key, {"name": "m"}) self.assertEqual(status, 201) pid = proj["id"] _insert_task(self.db_path, "m1", status="completed", project_id=pid, started=1000.0, finished=1002.0) _insert_task(self.db_path, "m2", status="failed", project_id=pid, started=1000.0, finished=1001.0) _insert_task(self.db_path, "other", status="completed", project_id="someone-else") costs = CostStore(self.db_path) costs.record("m1", pid, "gpt-6.1-sol", input_tokens=10, output_tokens=5, source="actual", fallback_count=1) costs.record("m2", pid, "gpt-6.1-sol", input_tokens=20, output_tokens=0, source="estimated") status, metrics = self._req("GET", "/metrics", self.owner_key) self.assertEqual(status, 200) self.assertEqual(metrics["total_tasks"], 2) self.assertEqual(metrics["by_status"]["completed"], 1) self.assertEqual(metrics["by_status"]["failed"], 1) self.assertAlmostEqual(metrics["success_rate"], 0.5) self.assertAlmostEqual(metrics["avg_duration_seconds"], 1.5) self.assertEqual(metrics["fallback_total"], 1) self.assertEqual(metrics["cost"]["actual_tokens"], 15) self.assertEqual(metrics["cost"]["estimated_tokens"], 20) # 新增观测键：队列深度、worker 聚合、分模型、未定价占比。 self.assertEqual(metrics["queue_depth"], {"queued": 0, "running": 0}) self.assertEqual(metrics["workers"]["total"], 0) self.assertEqual(metrics["failures_last_24h"], 0) models = {m["model"]: m for m in metrics["cost"]["by_model"]} self.assertEqual(models["gpt-6.1-sol"]["total_tokens"], 35) self.assertIsNotNone(metrics["cost"]["unpriced_share"]) # 无 key 401。 status, _ = self._req("GET", "/metrics") self.assertEqual(status, 401) def test_backup_rotation_keeps_three(self): init_db(self.db_path) _insert_task(self.db_path, "b1") out_dir = os.path.join(self.tmp.name, "backups") for stamp in ("20260101000000", "20260102000000", "20260103000000", "20260104000000", "20260105000000"): backup_db(self.db_path, out_dir, keep=3, timestamp=stamp) files = sorted(os.listdir(out_dir)) self.assertEqual(len(files), 3) self.assertEqual(files[0], "cloud-backup-20260103000000.db") conn = sqlite3.connect(os.path.join(out_dir, files[-1])) try: row = conn.execute( "SELECT id FROM tasks WHERE id = 'b1'").fetchone() self.assertIsNotNone(row) finally: conn.close() def test_postgres_backend_without_driver_fails_before_file(self): # 环境无关：用 mock 模拟驱动缺失，而不是依赖本机真的没装驱动 # （另一任务已在本 VM 装了 psycopg2 + Postgres，原写法在该环境下 # 前提不存在而失败，并非代码回归；CI runner 无驱动时原逻辑同样成立）。 other_db = os.path.join(self.tmp.name, "pg.db") with mock.patch("cloud.db._load_postgres_driver", side_effect=PostgresDriverMissingError("no driver")): with self.assertRaises(PostgresDriverMissingError): create_server(other_db, "127.0.0.1", 0, backend="postgres") self.assertFalse(os.path.exists(other_db)) class LoggingTestCase(unittest.TestCase): def test_json_log_contains_task_and_event_without_secret(self): stream = io.StringIO() handler = logging.StreamHandler(stream) handler.setFormatter(JsonFormatter()) logger = logging.getLogger("cloud.test-prod-gaps") logger.handlers = [handler] logger.setLevel(logging.INFO) logger.propagate = False try: log_event(logger, "worker_claim", task_id="t9", project_id="p9", status="running") finally: logger.handlers = [] payload = json.loads(stream.getvalue().strip()) self.assertEqual(payload["task_id"], "t9") self.assertEqual(payload["event"], "worker_claim") self.assertEqual(payload["service"], "cloud-control-plane") # setup_logging 本身不重复加 handler。 root = setup_logging(stream=io.StringIO()) self.assertEqual(len(root.handlers), 1) if __name__ == "__main__": unittest.main()