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()