356 lines
14 KiB
Python
356 lines
14 KiB
Python
"""生产级缺口单测：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()
|