440 lines
18 KiB
Python
440 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from datetime import datetime, timedelta
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from psycopg2 import errors
|
|
|
|
from app.services import cluster_runtime_service
|
|
|
|
|
|
class RuntimeSchemaGuardTests(unittest.TestCase):
|
|
def test_ensure_runtime_schema_does_not_toggle_autocommit_after_queries(self) -> None:
|
|
class _Connection:
|
|
def __init__(self) -> None:
|
|
self._autocommit = False
|
|
self.touched = False
|
|
self.committed = False
|
|
|
|
@property
|
|
def autocommit(self):
|
|
return self._autocommit
|
|
|
|
@autocommit.setter
|
|
def autocommit(self, value):
|
|
if self.touched:
|
|
raise AssertionError("autocommit should not be reassigned after queries start")
|
|
self._autocommit = value
|
|
|
|
def cursor(self):
|
|
conn = self
|
|
|
|
class _CursorContext:
|
|
def __enter__(self_inner):
|
|
conn.touched = True
|
|
return cursor
|
|
|
|
def __exit__(self_inner, exc_type, exc, tb):
|
|
return False
|
|
|
|
return _CursorContext()
|
|
|
|
def commit(self):
|
|
self.committed = True
|
|
|
|
cursor = MagicMock()
|
|
conn = _Connection()
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "_RUNTIME_SCHEMA_READY", False):
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm):
|
|
with patch.object(cluster_runtime_service, "_runtime_schema_basics_present", return_value=False):
|
|
cluster_runtime_service.ensure_runtime_schema()
|
|
|
|
cursor.execute.assert_any_call(
|
|
"SELECT pg_advisory_xact_lock(%s)",
|
|
(cluster_runtime_service._RUNTIME_SCHEMA_ADVISORY_LOCK_ID,),
|
|
)
|
|
cursor.execute.assert_any_call(cluster_runtime_service._RUNTIME_SCHEMA_SQL)
|
|
self.assertTrue(conn.committed)
|
|
|
|
def test_ensure_runtime_schema_executes_only_once_per_process(self) -> None:
|
|
conn = MagicMock()
|
|
cursor_cm = MagicMock()
|
|
cursor = MagicMock()
|
|
conn.cursor.return_value = cursor_cm
|
|
cursor_cm.__enter__.return_value = cursor
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "_RUNTIME_SCHEMA_READY", False):
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm) as mocked_get_db:
|
|
with patch.object(cluster_runtime_service, "_runtime_schema_basics_present", return_value=False):
|
|
cluster_runtime_service.ensure_runtime_schema()
|
|
cluster_runtime_service.ensure_runtime_schema()
|
|
|
|
mocked_get_db.assert_called_once()
|
|
cursor.execute.assert_any_call(
|
|
"SELECT pg_advisory_xact_lock(%s)",
|
|
(cluster_runtime_service._RUNTIME_SCHEMA_ADVISORY_LOCK_ID,),
|
|
)
|
|
cursor.execute.assert_any_call(cluster_runtime_service._RUNTIME_SCHEMA_SQL)
|
|
conn.commit.assert_called_once()
|
|
|
|
def test_ensure_runtime_schema_skips_ddl_when_required_schema_already_exists(self) -> None:
|
|
conn = MagicMock()
|
|
cursor_cm = MagicMock()
|
|
cursor = MagicMock()
|
|
conn.cursor.return_value = cursor_cm
|
|
cursor_cm.__enter__.return_value = cursor
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "_RUNTIME_SCHEMA_READY", False):
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm):
|
|
with patch.object(cluster_runtime_service, "_runtime_schema_basics_present", return_value=True):
|
|
with patch.object(cluster_runtime_service, "_runtime_missing_indexes", return_value=iter(())):
|
|
cluster_runtime_service.ensure_runtime_schema()
|
|
|
|
self.assertFalse(
|
|
any(sql == cluster_runtime_service._RUNTIME_SCHEMA_SQL for sql, _params in cursor.execute.call_args_list)
|
|
)
|
|
conn.commit.assert_not_called()
|
|
|
|
def test_ensure_runtime_schema_accepts_deadlock_when_required_schema_already_exists(self) -> None:
|
|
class _Cursor:
|
|
def __init__(self, *, raise_on_schema=False, fetchone_values=None, fetchall_values=None) -> None:
|
|
self.raise_on_schema = raise_on_schema
|
|
self.fetchone_values = list(fetchone_values or [])
|
|
self.fetchall_values = list(fetchall_values or [])
|
|
self.executed = []
|
|
|
|
def execute(self, sql, params=None):
|
|
self.executed.append((sql, params))
|
|
if self.raise_on_schema and sql == cluster_runtime_service._RUNTIME_SCHEMA_SQL:
|
|
raise errors.DeadlockDetected()
|
|
|
|
def fetchone(self):
|
|
if self.fetchone_values:
|
|
return self.fetchone_values.pop(0)
|
|
return None
|
|
|
|
def fetchall(self):
|
|
if self.fetchall_values:
|
|
return self.fetchall_values.pop(0)
|
|
return []
|
|
|
|
class _CursorContext:
|
|
def __init__(self, cursor) -> None:
|
|
self.cursor = cursor
|
|
|
|
def __enter__(self):
|
|
return self.cursor
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
conn = MagicMock()
|
|
first_cursor = _Cursor()
|
|
second_cursor = _Cursor(raise_on_schema=True)
|
|
third_cursor = _Cursor()
|
|
conn.cursor.side_effect = [
|
|
_CursorContext(first_cursor),
|
|
_CursorContext(second_cursor),
|
|
_CursorContext(third_cursor),
|
|
]
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "_RUNTIME_SCHEMA_READY", False):
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm):
|
|
with patch.object(cluster_runtime_service, "_runtime_schema_basics_present", side_effect=[False, True]):
|
|
with patch.object(cluster_runtime_service, "_runtime_missing_indexes", return_value=iter(())):
|
|
cluster_runtime_service.ensure_runtime_schema()
|
|
|
|
conn.rollback.assert_called_once()
|
|
conn.commit.assert_not_called()
|
|
|
|
def test_ensure_runtime_schema_repairs_missing_indexes_without_full_ddl(self) -> None:
|
|
conn = MagicMock()
|
|
cursor_cm = MagicMock()
|
|
cursor = MagicMock()
|
|
conn.cursor.return_value = cursor_cm
|
|
cursor_cm.__enter__.return_value = cursor
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "_RUNTIME_SCHEMA_READY", False):
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm):
|
|
with patch.object(cluster_runtime_service, "_runtime_schema_basics_present", return_value=True):
|
|
with patch.object(
|
|
cluster_runtime_service,
|
|
"_runtime_missing_indexes",
|
|
side_effect=[iter(("idx_detect_job_items_claim_step_ready",)), iter(())],
|
|
):
|
|
with patch.object(cluster_runtime_service, "_ensure_runtime_schema_indexes") as mocked_ensure_indexes:
|
|
cluster_runtime_service.ensure_runtime_schema()
|
|
|
|
mocked_ensure_indexes.assert_called_once()
|
|
self.assertFalse(
|
|
any(sql == cluster_runtime_service._RUNTIME_SCHEMA_SQL for sql, _params in cursor.execute.call_args_list)
|
|
)
|
|
|
|
def test_runtime_missing_indexes_treats_invalid_indexes_as_missing(self) -> None:
|
|
cur = MagicMock()
|
|
cur.fetchall.return_value = [
|
|
("idx_detect_job_items_job_domain_step", True, True, True),
|
|
("idx_detect_job_items_claim_step_ready", False, True, True),
|
|
]
|
|
|
|
missing = list(cluster_runtime_service._runtime_missing_indexes(cur))
|
|
|
|
self.assertIn("idx_detect_job_items_claim_step_ready", missing)
|
|
self.assertIn("idx_detect_job_items_claim_job_step_ready", missing)
|
|
self.assertNotIn("idx_detect_job_items_job_domain_step", missing)
|
|
|
|
def test_control_node_supports_worker_only_on_mainland_with_worker_signals(self) -> None:
|
|
self.assertFalse(
|
|
cluster_runtime_service._control_node_supports_worker(
|
|
region="overseas",
|
|
metadata={
|
|
"worker_online": False,
|
|
"detect_participating": True,
|
|
"active_threads": 0,
|
|
"max_threads": 0,
|
|
},
|
|
)
|
|
)
|
|
self.assertTrue(
|
|
cluster_runtime_service._control_node_supports_worker(
|
|
region="mainland",
|
|
metadata={
|
|
"worker_online": True,
|
|
"detect_participating": False,
|
|
"active_threads": 0,
|
|
"max_threads": 0,
|
|
},
|
|
)
|
|
)
|
|
|
|
def test_metadata_idle_without_runtime_work_detects_stale_idle_heartbeat(self) -> None:
|
|
self.assertTrue(
|
|
cluster_runtime_service._metadata_idle_without_runtime_work(
|
|
{
|
|
"phase_label": "idle",
|
|
"phase_detail": "Worker 已启动，等待检测指令",
|
|
"active_threads": 323,
|
|
"max_threads": 4,
|
|
"active_job_code": "",
|
|
"job_items_total": 0,
|
|
"job_items_claimed": 0,
|
|
"job_items_running": 0,
|
|
"job_items_completed": 0,
|
|
"job_items_failed": 0,
|
|
}
|
|
)
|
|
)
|
|
self.assertFalse(
|
|
cluster_runtime_service._metadata_idle_without_runtime_work(
|
|
{
|
|
"phase_label": "running",
|
|
"active_threads": 12,
|
|
}
|
|
)
|
|
)
|
|
self.assertFalse(
|
|
cluster_runtime_service._metadata_idle_without_runtime_work(
|
|
{
|
|
"phase_label": "idle",
|
|
"active_job_code": "sync-overseas-1",
|
|
}
|
|
)
|
|
)
|
|
|
|
def test_get_cluster_snapshot_counts_all_rows_even_when_display_nodes_are_limited(self) -> None:
|
|
now = datetime.now()
|
|
display_rows = [
|
|
(
|
|
f"mainland-worker-{index:03d}",
|
|
"mainland",
|
|
"worker",
|
|
f"worker-{index:03d}",
|
|
f"10.0.0.{index}",
|
|
"online",
|
|
"test",
|
|
0,
|
|
{},
|
|
now,
|
|
now,
|
|
)
|
|
for index in range(1, 101)
|
|
]
|
|
summary_rows = display_rows + [
|
|
(
|
|
"mainland-worker-101",
|
|
"mainland",
|
|
"worker",
|
|
"worker-101",
|
|
"10.0.0.101",
|
|
"online",
|
|
"test",
|
|
0,
|
|
{},
|
|
now,
|
|
now,
|
|
)
|
|
]
|
|
conn = MagicMock()
|
|
cursor_cm = MagicMock()
|
|
cursor = MagicMock()
|
|
cursor.fetchall.side_effect = [display_rows, summary_rows]
|
|
cursor.fetchone.side_effect = [(12,), (34,)]
|
|
conn.cursor.return_value = cursor_cm
|
|
cursor_cm.__enter__.return_value = cursor
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm):
|
|
with patch.object(cluster_runtime_service, "prune_expired_runtime_nodes"):
|
|
with patch.object(cluster_runtime_service, "register_local_control_heartbeat"):
|
|
with patch.object(cluster_runtime_service, "_load_managed_node_overlays", return_value={}):
|
|
with patch.object(cluster_runtime_service, "_load_disabled_managed_node_codes", return_value=set()):
|
|
snapshot = cluster_runtime_service.get_cluster_snapshot()
|
|
|
|
self.assertEqual(100, len(snapshot["nodes"]))
|
|
self.assertEqual(101, snapshot["nodes_total"])
|
|
self.assertEqual(101, snapshot["summary"]["online_worker_nodes"])
|
|
self.assertEqual(12, snapshot["jobs_total"])
|
|
self.assertEqual(34, snapshot["active_job_items"])
|
|
|
|
def test_get_cluster_snapshot_excludes_disabled_managed_nodes(self) -> None:
|
|
now = datetime.now()
|
|
display_rows = [
|
|
("mainland-controller-01", "mainland", "control", "controller", "10.0.0.1", "online", "test", 1, {}, now, now),
|
|
("mainland-worker-01", "mainland", "worker", "worker-01", "10.0.0.2", "busy", "test", 1, {}, now, now),
|
|
]
|
|
summary_rows = list(display_rows)
|
|
conn = MagicMock()
|
|
cursor_cm = MagicMock()
|
|
cursor = MagicMock()
|
|
cursor.fetchall.side_effect = [display_rows, summary_rows]
|
|
cursor.fetchone.side_effect = [(1,), (2,)]
|
|
conn.cursor.return_value = cursor_cm
|
|
cursor_cm.__enter__.return_value = cursor
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm):
|
|
with patch.object(cluster_runtime_service, "prune_expired_runtime_nodes"):
|
|
with patch.object(cluster_runtime_service, "register_local_control_heartbeat"):
|
|
with patch.object(cluster_runtime_service, "_load_managed_node_overlays", return_value={}):
|
|
with patch.object(cluster_runtime_service, "_load_disabled_managed_node_codes", return_value={"mainland-worker-01"}):
|
|
snapshot = cluster_runtime_service.get_cluster_snapshot()
|
|
|
|
self.assertEqual(["mainland-controller-01"], [item["node_code"] for item in snapshot["nodes"]])
|
|
self.assertEqual(1, snapshot["nodes_total"])
|
|
self.assertEqual(1, snapshot["summary"]["online_control_nodes"])
|
|
self.assertEqual(0, snapshot["summary"]["online_worker_nodes"])
|
|
|
|
def test_get_cluster_snapshot_prefers_imported_runtime_update_time_for_freshness(self) -> None:
|
|
now = datetime.now()
|
|
stale_heartbeat = now - timedelta(minutes=8)
|
|
fresh_ingest = now - timedelta(seconds=20)
|
|
display_rows = [
|
|
(
|
|
"mainland-controller-01-da",
|
|
"mainland",
|
|
"worker",
|
|
"controller",
|
|
"10.0.0.1",
|
|
"busy",
|
|
"test",
|
|
42,
|
|
{
|
|
"service": "runtime-ingest",
|
|
"updated_at": fresh_ingest.isoformat(timespec="seconds"),
|
|
"active_threads": 42,
|
|
},
|
|
stale_heartbeat,
|
|
fresh_ingest,
|
|
)
|
|
]
|
|
summary_rows = list(display_rows)
|
|
conn = MagicMock()
|
|
cursor_cm = MagicMock()
|
|
cursor = MagicMock()
|
|
cursor.fetchall.side_effect = [display_rows, summary_rows]
|
|
cursor.fetchone.side_effect = [(1,), (1,)]
|
|
conn.cursor.return_value = cursor_cm
|
|
cursor_cm.__enter__.return_value = cursor
|
|
db_cm = MagicMock()
|
|
db_cm.__enter__.return_value = conn
|
|
|
|
with patch.object(cluster_runtime_service, "get_db", return_value=db_cm):
|
|
with patch.object(cluster_runtime_service, "prune_expired_runtime_nodes"):
|
|
with patch.object(cluster_runtime_service, "register_local_control_heartbeat"):
|
|
with patch.object(cluster_runtime_service, "_load_managed_node_overlays", return_value={}):
|
|
with patch.object(cluster_runtime_service, "_load_disabled_managed_node_codes", return_value=set()):
|
|
snapshot = cluster_runtime_service.get_cluster_snapshot()
|
|
|
|
self.assertEqual(1, snapshot["summary"]["online_worker_nodes"])
|
|
self.assertEqual([], snapshot["summary"]["offline_nodes"])
|
|
self.assertEqual("busy", snapshot["nodes"][0]["status"])
|
|
|
|
def test_register_local_control_heartbeat_uses_detect_status_snapshot(self) -> None:
|
|
detect_status = {
|
|
"worker_online": True,
|
|
"active_thread_count": 17,
|
|
"max_thread_count": 320,
|
|
"available_proxy_count": 41,
|
|
"proxy_runtime_label": "代理可用",
|
|
"proxy_runtime_reason": "pool_ready",
|
|
"proxy_last_refresh_status": "ok",
|
|
"proxy_last_refresh_time": "2026-04-24 01:30:00",
|
|
"proxy_last_refresh_source_count": 3,
|
|
"proxy_last_refresh_total_items": 120,
|
|
"proxy_last_validated_count": 110,
|
|
"proxy_last_available_count": 41,
|
|
"phase_label": "running",
|
|
"phase_detail": "正在执行检测",
|
|
"active_job": {
|
|
"job_code": "detect-20260424013000-abcd12",
|
|
"status": "running",
|
|
"node_stats": [
|
|
{
|
|
"node_code": "mainland-controller-01",
|
|
"items_total": 56,
|
|
"items_claimed": 5,
|
|
"items_running": 9,
|
|
"items_completed": 42,
|
|
}
|
|
],
|
|
},
|
|
}
|
|
|
|
with patch.object(cluster_runtime_service.settings, "node_code", "mainland-controller-01"):
|
|
with patch.object(cluster_runtime_service.settings, "node_region", "mainland"):
|
|
with patch.object(cluster_runtime_service.settings, "node_role", "control"):
|
|
with patch("app.services.detect_service.get_detect_status", return_value=detect_status):
|
|
with patch.object(cluster_runtime_service, "register_node_heartbeat") as mock_register:
|
|
cluster_runtime_service.register_local_control_heartbeat()
|
|
|
|
mock_register.assert_called_once()
|
|
payload = mock_register.call_args.kwargs
|
|
self.assertEqual("mainland-controller-01", payload["node_code"])
|
|
self.assertEqual("busy", payload["status"])
|
|
self.assertEqual(17, payload["current_load"])
|
|
self.assertTrue(payload["metadata"]["worker_online"])
|
|
self.assertTrue(payload["metadata"]["detect_participating"])
|
|
self.assertEqual("detect-20260424013000-abcd12", payload["metadata"]["active_job_code"])
|
|
self.assertEqual(17, payload["metadata"]["active_threads"])
|
|
self.assertEqual(320, payload["metadata"]["max_threads"])
|
|
self.assertEqual(41, payload["metadata"]["available_proxy_count"])
|
|
self.assertEqual("running", payload["metadata"]["phase_label"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|