238 lines
9.4 KiB
Python
238 lines
9.4 KiB
Python
"""运维批单测：恢复演练、worker 心跳、health worker 字段、monitor 自测。
|
|
|
|
只用标准库，不跑真实模型回合、不联网、不连真实 GitHub、不连生产库；
|
|
restore 的 psql 与 monitor 的 systemctl/docker 一律用假 runner 或
|
|
--self-test 的 canned server 验证，不在本机起真实服务。
|
|
"""
|
|
|
|
import http.client
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
try:
|
|
from cloud.backup import restore_postgres
|
|
from cloud.server import create_server
|
|
from cloud.worker import Worker
|
|
from cloud.workers import heartbeat_worker, register_worker
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.backup import restore_postgres
|
|
from cloud.server import create_server
|
|
from cloud.worker import Worker
|
|
from cloud.workers import heartbeat_worker, register_worker
|
|
|
|
|
|
def _fake_result(returncode=0, stdout="", stderr=""):
|
|
return SimpleNamespace(returncode=returncode, stdout=stdout,
|
|
stderr=stderr)
|
|
|
|
|
|
class RestorePostgresTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.backup = os.path.join(self.tmp.name, "backup.sql")
|
|
with open(self.backup, "w") as f:
|
|
f.write("-- fake dump\n")
|
|
self.old_env = dict(os.environ)
|
|
os.environ["CLOUD_DB_HOST"] = "127.0.0.1"
|
|
os.environ["CLOUD_DB_PORT"] = "5432"
|
|
os.environ["CLOUD_DB_NAME"] = "prod_db"
|
|
os.environ["CLOUD_DB_USER"] = "codexsvc"
|
|
os.environ["CLOUD_DB_PASSWORD"] = "test-secret-pw"
|
|
|
|
def tearDown(self):
|
|
os.environ.clear()
|
|
os.environ.update(self.old_env)
|
|
self.tmp.cleanup()
|
|
|
|
def test_target_equal_to_production_rejected(self):
|
|
# 目标库名等于生产库名时直接拒绝，allow_production 默认 False。
|
|
with self.assertRaises(ValueError):
|
|
restore_postgres(self.backup, "prod_db",
|
|
runner=lambda c, e: _fake_result())
|
|
|
|
def test_target_equal_to_production_allowed_only_explicitly(self):
|
|
calls = []
|
|
|
|
def runner(cmd, env):
|
|
calls.append((list(cmd), dict(env)))
|
|
return _fake_result()
|
|
|
|
out = restore_postgres(self.backup, "prod_db", runner=runner,
|
|
allow_production=True)
|
|
self.assertEqual(out, "prod_db")
|
|
self.assertEqual(len(calls), 1)
|
|
|
|
def test_missing_backup_file_rejected(self):
|
|
with self.assertRaises(FileNotFoundError):
|
|
restore_postgres(os.path.join(self.tmp.name, "nope.sql"),
|
|
"restore_test",
|
|
runner=lambda c, e: _fake_result())
|
|
|
|
def test_empty_backup_file_rejected(self):
|
|
empty = os.path.join(self.tmp.name, "empty.sql")
|
|
open(empty, "w").close()
|
|
with self.assertRaises(ValueError):
|
|
restore_postgres(empty, "restore_test",
|
|
runner=lambda c, e: _fake_result())
|
|
|
|
def test_password_only_in_env_not_in_command(self):
|
|
calls = []
|
|
|
|
def runner(cmd, env):
|
|
calls.append((list(cmd), dict(env)))
|
|
return _fake_result()
|
|
|
|
out = restore_postgres(self.backup, "restore_test", runner=runner)
|
|
self.assertEqual(out, "restore_test")
|
|
cmd, env = calls[0]
|
|
self.assertNotIn("test-secret-pw", " ".join(cmd))
|
|
self.assertEqual(env.get("PGPASSWORD"), "test-secret-pw")
|
|
self.assertIn("-f", cmd)
|
|
self.assertIn(self.backup, cmd)
|
|
self.assertIn("restore_test", cmd)
|
|
|
|
def test_psql_failure_raises(self):
|
|
def runner(cmd, env):
|
|
return _fake_result(returncode=1, stderr="psql boom")
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
restore_postgres(self.backup, "restore_test", runner=runner)
|
|
|
|
|
|
class WorkerHeartbeatTestCase(unittest.TestCase):
|
|
def test_heartbeat_file_written_with_last_seen_and_pid(self):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
hb = os.path.join(tmp.name, "worker-heartbeat.json")
|
|
worker = Worker(os.path.join(tmp.name, "db.db"), "/nonexistent/host.py",
|
|
env_key="TEST_KEY", heartbeat_file=hb)
|
|
worker._write_heartbeat()
|
|
with open(hb) as f:
|
|
payload = json.load(f)
|
|
self.assertIn("last_seen", payload)
|
|
self.assertEqual(payload["pid"], os.getpid())
|
|
self.assertLess(abs(payload["last_seen"] - time.time()), 5)
|
|
tmp.cleanup()
|
|
|
|
def test_heartbeat_write_failure_does_not_crash_worker(self):
|
|
# 心跳目录不可写时只记日志，不抛异常（写到一个文件路径下当目录）。
|
|
tmp = tempfile.TemporaryDirectory()
|
|
blocker = os.path.join(tmp.name, "blocker")
|
|
open(blocker, "w").close()
|
|
worker = Worker(os.path.join(tmp.name, "db.db"), "/nonexistent/host.py",
|
|
env_key="TEST_KEY",
|
|
heartbeat_file=os.path.join(blocker, "hb.json"))
|
|
worker._write_heartbeat() # 不抛异常即通过
|
|
tmp.cleanup()
|
|
|
|
|
|
class HealthWorkerFieldTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
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 _health(self):
|
|
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
|
conn.request("GET", "/health")
|
|
resp = conn.getresponse()
|
|
body = json.loads(resp.read())
|
|
conn.close()
|
|
return resp.status, body
|
|
|
|
def test_health_worker_fresh_ok(self):
|
|
hb = os.path.join(self.tmp.name, "hb.json")
|
|
with open(hb, "w") as f:
|
|
json.dump({"last_seen": time.time(), "pid": 123}, f)
|
|
self.server.worker_heartbeat_file = hb
|
|
status, body = self._health()
|
|
self.assertEqual(status, 200)
|
|
self.assertTrue(body["ok"])
|
|
self.assertTrue(body["worker_ok"])
|
|
self.assertIsNotNone(body["worker_last_seen_seconds_ago"])
|
|
self.assertLess(body["worker_last_seen_seconds_ago"], 5)
|
|
|
|
def test_health_worker_stale_not_ok_but_still_200(self):
|
|
hb = os.path.join(self.tmp.name, "hb.json")
|
|
with open(hb, "w") as f:
|
|
json.dump({"last_seen": time.time() - 1000, "pid": 123}, f)
|
|
self.server.worker_heartbeat_file = hb
|
|
status, body = self._health()
|
|
# worker 不 ok 不代表 server 挂了，health 本身仍返回 200。
|
|
self.assertEqual(status, 200)
|
|
self.assertFalse(body["worker_ok"])
|
|
self.assertGreater(body["worker_last_seen_seconds_ago"], 500)
|
|
|
|
def test_health_worker_missing_file_none_false(self):
|
|
self.server.worker_heartbeat_file = os.path.join(
|
|
self.tmp.name, "no-such-hb.json")
|
|
status, body = self._health()
|
|
self.assertEqual(status, 200)
|
|
self.assertIsNone(body["worker_last_seen_seconds_ago"])
|
|
self.assertFalse(body["worker_ok"])
|
|
|
|
def test_health_worker_db_registry_fresh_ok_without_file(self):
|
|
# 多机前提：文件心跳缺失时，DB 注册表的新鲜心跳同样算 ok。
|
|
register_worker(self.db_path, "w-db", executor="docker", capacity=1)
|
|
heartbeat_worker(self.db_path, "w-db", executor="docker")
|
|
self.server.worker_heartbeat_file = os.path.join(
|
|
self.tmp.name, "no-such-hb.json")
|
|
status, body = self._health()
|
|
self.assertEqual(status, 200)
|
|
self.assertTrue(body["worker_ok"])
|
|
self.assertIsNotNone(body["worker_last_seen_seconds_ago"])
|
|
self.assertLess(body["worker_last_seen_seconds_ago"], 5)
|
|
|
|
def test_health_worker_db_registry_stale_not_ok(self):
|
|
register_worker(self.db_path, "w-db", executor="docker", capacity=1)
|
|
import sqlite3
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute("UPDATE workers SET last_seen = ?",
|
|
(time.time() - 1000,))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
self.server.worker_heartbeat_file = os.path.join(
|
|
self.tmp.name, "no-such-hb.json")
|
|
status, body = self._health()
|
|
self.assertEqual(status, 200)
|
|
self.assertFalse(body["worker_ok"])
|
|
self.assertGreater(body["worker_last_seen_seconds_ago"], 500)
|
|
|
|
|
|
class MonitorSelfTestCase(unittest.TestCase):
|
|
def test_monitor_self_test_passes(self):
|
|
# monitor.sh --self-test 只验证脚本机制（health 解析 + JSON 合法），
|
|
# 不断言主机健康；VM 上其它项 fail 是正常的，self-test 仍应通过。
|
|
repo_root = Path(__file__).resolve().parents[2]
|
|
script = repo_root / "deploy" / "monitor.sh"
|
|
self.assertTrue(script.exists())
|
|
result = subprocess.run(
|
|
["sh", str(script), "--self-test"],
|
|
capture_output=True, text=True, timeout=60)
|
|
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
|
|
self.assertIn("SELF-TEST PASS", result.stdout)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|