"""workers 表 host/pid/address 落地单测：幂等迁移 + 心跳 + 失联判定。 只用标准库 + 本地 SQLite 临时库，不跑真实模型、不联网、不杀 任何进程或容器：失联判定只测 classify_stale_workers 的返回值， pid 存活一律用注入的 pid_is_dead_fn，不依赖本机真实进程号； pid_is_dead 本体只测非法 pid（False）与不存在的大进程号 （True），不测当前存活进程（环境相关）。 """ import os import sqlite3 import sys import tempfile import time import unittest from pathlib import Path try: from cloud.workers import ( classify_stale_workers, ensure_workers_table, heartbeat_worker, list_workers, pid_is_dead, register_worker, ) except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.workers import ( classify_stale_workers, ensure_workers_table, heartbeat_worker, list_workers, pid_is_dead, register_worker, ) def _cols(db_path): conn = sqlite3.connect(db_path) try: rows = conn.execute("PRAGMA table_info(workers)").fetchall() return {row[1]: row[2] for row in rows} finally: conn.close() def _set_last_seen(db_path, worker_id, last_seen): conn = sqlite3.connect(db_path) try: conn.execute("UPDATE workers SET last_seen = ? WHERE worker_id = ?", (last_seen, worker_id)) conn.commit() finally: conn.close() def _by_id(rows, worker_id): for row in rows: if row["worker_id"] == worker_id: return row raise AssertionError("missing worker %s" % worker_id) class HostpidMigrationTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") def tearDown(self): self.tmp.cleanup() def test_new_db_table_has_host_pid_address(self): ensure_workers_table(self.db_path) cols = _cols(self.db_path) self.assertIn("host", cols) self.assertIn("pid", cols) self.assertIn("address", cols) def test_migration_is_idempotent(self): ensure_workers_table(self.db_path) register_worker(self.db_path, "w1", executor="local", host="h1", pid=1234, address="10.0.0.1") ensure_workers_table(self.db_path) ensure_workers_table(self.db_path) rows = list_workers(self.db_path) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["host"], "h1") self.assertEqual(rows[0]["pid"], 1234) self.assertEqual(rows[0]["address"], "10.0.0.1") def test_old_db_upgrade_keeps_rows_and_leaves_new_cols_null(self): # 手工建旧表（无 host/pid/address 列、无 running 列外的 # 旧列集），写一行，再跑迁移：旧行保留、新列为空。 conn = sqlite3.connect(self.db_path) try: conn.execute( "CREATE TABLE workers(" "worker_id TEXT PRIMARY KEY, executor TEXT, " "capacity INTEGER, last_seen DOUBLE PRECISION, " "created_at DOUBLE PRECISION)") conn.execute( "INSERT INTO workers VALUES('w-old', 'docker', 1, ?, ?)", (1000.0, 900.0)) conn.commit() finally: conn.close() ensure_workers_table(self.db_path) ensure_workers_table(self.db_path) rows = list_workers(self.db_path) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["executor"], "docker") self.assertEqual(rows[0]["last_seen"], 1000.0) self.assertIsNone(rows[0]["host"]) self.assertIsNone(rows[0]["pid"]) self.assertIsNone(rows[0]["address"]) def test_heartbeat_refreshes_host_pid_and_preserves_when_omitted(self): register_worker(self.db_path, "w1", executor="local", host="h-old", pid=111, address="a-old") row = heartbeat_worker(self.db_path, "w1", host="h-new", pid=222) self.assertEqual(row["host"], "h-new") self.assertEqual(row["pid"], 222) # address 未给：保留旧值，不被洗成 None。 self.assertEqual(row["address"], "a-old") row = heartbeat_worker(self.db_path, "w1") self.assertEqual(row["host"], "h-new") self.assertEqual(row["pid"], 222) self.assertEqual(row["address"], "a-old") def test_heartbeat_new_row_stores_given_values_and_nulls_when_omitted(self): row = heartbeat_worker(self.db_path, "w-new", host="h1", pid=333) self.assertEqual(row["host"], "h1") self.assertEqual(row["pid"], 333) self.assertIsNone(row["address"]) row = heartbeat_worker(self.db_path, "w-plain") self.assertIsNone(row["host"]) self.assertIsNone(row["pid"]) self.assertIsNone(row["address"]) def test_invalid_pid_rejected(self): with self.assertRaises(ValueError): heartbeat_worker(self.db_path, "w1", pid=0) with self.assertRaises(ValueError): heartbeat_worker(self.db_path, "w1", pid="not-a-pid") with self.assertRaises(ValueError): heartbeat_worker(self.db_path, "w1", pid=True) class StaleClassificationTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") self.now = time.time() # 阈值 120 秒：w-kill 超阈 + 本机 + 已死；w-eq 恰等于阈值； # w-foreign 超阈但异机；w-nopid 超阈但无 pid；w-alive 超阈 # 但 pid 还活着；w-fresh 未超阈。 register_worker(self.db_path, "w-kill", host="h-local", pid=9001) register_worker(self.db_path, "w-eq", host="h-local", pid=9002) register_worker(self.db_path, "w-foreign", host="h-other", pid=9003) register_worker(self.db_path, "w-nopid", host="h-local") register_worker(self.db_path, "w-alive", host="h-local", pid=9004) register_worker(self.db_path, "w-fresh", host="h-local", pid=9005) _set_last_seen(self.db_path, "w-kill", self.now - 121) _set_last_seen(self.db_path, "w-eq", self.now - 120) _set_last_seen(self.db_path, "w-foreign", self.now - 200) _set_last_seen(self.db_path, "w-nopid", self.now - 200) _set_last_seen(self.db_path, "w-alive", self.now - 200) _set_last_seen(self.db_path, "w-fresh", self.now - 10) def tearDown(self): self.tmp.cleanup() def _classify(self): dead = {9001, 9002, 9003} return classify_stale_workers( self.db_path, stale_seconds=120, local_host="h-local", pid_is_dead_fn=lambda pid: pid in dead, now=self.now) def test_only_pid_verified_stale_is_kill(self): rows = self._classify() self.assertEqual(_by_id(rows, "w-kill")["action"], "kill") self.assertTrue(_by_id(rows, "w-kill")["pid_verified"]) self.assertEqual(_by_id(rows, "w-kill")["reason"], "pid_verified_dead") def test_equal_threshold_is_list_only(self): rows = self._classify() row = _by_id(rows, "w-eq") self.assertEqual(row["action"], "list_only") self.assertFalse(row["pid_verified"]) self.assertEqual(row["reason"], "age_equals_threshold") def test_non_verified_stale_are_list_only(self): rows = self._classify() self.assertEqual(_by_id(rows, "w-foreign")["action"], "list_only") self.assertEqual(_by_id(rows, "w-foreign")["reason"], "host_mismatch") self.assertEqual(_by_id(rows, "w-nopid")["action"], "list_only") self.assertEqual(_by_id(rows, "w-nopid")["reason"], "pid_missing") self.assertEqual(_by_id(rows, "w-alive")["action"], "list_only") self.assertEqual(_by_id(rows, "w-alive")["reason"], "pid_alive_or_unknown") self.assertEqual(_by_id(rows, "w-fresh")["action"], "list_only") self.assertEqual(_by_id(rows, "w-fresh")["reason"], "not_stale") def test_classification_kills_nothing(self): # 判定函数本身绝不杀进程：跑完后本进程仍在，且返回值里 # 除 w-kill 外无 kill 项；函数不接受任何容器名/前缀参数， # 无法按名删容器（签名即证据，见 classify_stale_workers）。 rows = self._classify() kills = [r for r in rows if r["action"] == "kill"] self.assertEqual([r["worker_id"] for r in kills], ["w-kill"]) self.assertTrue(pid_is_dead(999999)) class PgBulkSchemaBoundaryTestCase(unittest.TestCase): def test_workers_stays_out_of_postgres_bulk_schema(self): # 界线钉住：workers 是易失心跳数据（migrations.py 的 # UNMIGRATED_TABLES），PG 批量 schema 刻意不含它；PG # 生产库的 workers 表由 ensure_workers_table 懒建/补列。 # 本轨曾把 workers 加进 POSTGRES_SCHEMA，致 P2 迁移对齐 # 3 项变红（迁移表集、未迁移清单、覆盖率），已回退。 try: from cloud.db import POSTGRES_SCHEMA except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.db import POSTGRES_SCHEMA joined = "\n".join(POSTGRES_SCHEMA) self.assertNotIn("CREATE TABLE IF NOT EXISTS workers(", joined) class PidIsDeadTestCase(unittest.TestCase): def test_invalid_pid_is_not_dead(self): self.assertFalse(pid_is_dead(None)) self.assertFalse(pid_is_dead(0)) self.assertFalse(pid_is_dead("x")) if __name__ == "__main__": unittest.main()