cloud: production gaps (worker loop + true cancel, task result columns + migration, DB interface migration, metrics + backup)
This commit is contained in:
host committed
1 parent
07ef869267
commit
631580beab
14 files changed
+956
-29
No files matched your search
@@ -0,0 +1,342 @@
|
||||
"""生产级缺口单测：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 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)
|
||||
# 无 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):
|
||||
other_db = os.path.join(self.tmp.name, "pg.db")
|
||||
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()
|
||||
Reference in new issue
Block a user