d

domainCheck/tests/test_detect_worker_flush_retry.py

@@ -0,0 +1,225 @@
+import collections
+import threading
+import unittest
+import sys
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+
+sys.path.insert(0, "/www/wwwroot/getDomain/domainCheck")
+
+from detect_worker import DetectWorker  # noqa: E402
+
+
+class _FakeDb:
+    def __init__(self, method_name: str):
+        self.method_name = method_name
+        self.calls = 0
+
+    def _run(self, batch):
+        self.calls += 1
+        return -1 if self.calls == 1 else len(batch)
+
+    def mark_detect_job_items_running_batch(self, batch):
+        return self._run(batch)
+
+    def finalize_detect_job_items_batch(self, batch):
+        return self._run(batch)
+
+
+class DetectWorkerFlushRetryTest(unittest.TestCase):
+    def _build_worker(self, db):
+        worker = DetectWorker.__new__(DetectWorker)
+        worker.db = db
+        worker._task_local = threading.local()
+        worker._running_mark_lock = threading.Lock()
+        worker._pending_running_marks = collections.deque()
+        worker._last_running_mark_flush_at = 0.0
+        worker._job_finalize_lock = threading.Lock()
+        worker._pending_job_finalizations = collections.deque()
+        worker._last_job_finalize_flush_at = 0.0
+        worker.proxy_pool = []
+        worker.proxy_pool_lock = threading.Lock()
+        worker.proxy_failure_lock = threading.Lock()
+        worker.proxy_failure_counts = {}
+        worker.proxy_quarantine_until = {}
+        worker.proxy_source_by_key = {}
+        worker.proxy_source_failure_counts = {}
+        worker.proxy_source_quarantine_until = {}
+        worker.proxy_config = {"proxy_enable": True, "allow_direct": True}
+        worker.proxy_last_refresh_status = "test"
+        worker.trigger_proxy_refresh = MagicMock()
+        worker.runtime_settings = {}
+        worker.thread_count = 600
+        worker.detecting = True
+        worker._get_active_domain_threads = MagicMock(return_value=0)
+        return worker
+
+    def test_running_mark_flush_requeues_batch_when_db_write_fails(self):
+        worker = self._build_worker(_FakeDb("running"))
+        worker._pending_running_marks.extend([(101, "token-a"), (102, "token-b")])
+
+        updated = worker._flush_pending_running_marks(force=True, batch_limit=10)
+        self.assertEqual(0, updated)
+        self.assertEqual([(101, "token-a"), (102, "token-b")], list(worker._pending_running_marks))
+
+        updated = worker._flush_pending_running_marks(force=True, batch_limit=10)
+        self.assertEqual(2, updated)
+        self.assertEqual([], list(worker._pending_running_marks))
+
+    def test_job_finalization_flush_requeues_batch_when_db_write_fails(self):
+        worker = self._build_worker(_FakeDb("finalize"))
+        worker._pending_job_finalizations.extend(
+            [
+                {"job_item_id": 201, "claim_token": "token-a", "final_status": "completed"},
+                {"job_item_id": 202, "claim_token": "token-b", "final_status": "failed"},
+            ]
+        )
+
+        updated = worker._flush_pending_job_finalizations(force=True, batch_limit=10)
+        self.assertEqual(0, updated)
+        self.assertEqual(2, len(worker._pending_job_finalizations))
+
+        updated = worker._flush_pending_job_finalizations(force=True, batch_limit=10)
+        self.assertEqual(2, updated)
+        self.assertEqual(0, len(worker._pending_job_finalizations))
+
+    def test_proxy_auth_failure_quarantines_whole_source_group(self):
+        worker = self._build_worker(_FakeDb("noop"))
+        bad_proxy = {"http": "http://bad-proxy", "https": "http://bad-proxy"}
+        good_proxy = {"http": "http://good-proxy", "https": "http://good-proxy"}
+        worker.proxy_pool = [
+            {"proxy": bad_proxy, "usage_count": 0, "source_tag": "D"},
+            {"proxy": good_proxy, "usage_count": 0, "source_tag": "E"},
+        ]
+        worker.proxy_source_by_key[worker._proxy_key(bad_proxy)] = "D"
+        worker.proxy_source_by_key[worker._proxy_key(good_proxy)] = "E"
+
+        source_tag, failure_count, cooldown_seconds, removed = worker._mark_proxy_source_failure(
+            bad_proxy,
+            "Tunnel connection failed: 407 Proxy Authentication Required",
+        )
+
+        self.assertEqual("D", source_tag)
+        self.assertEqual(1, failure_count)
+        self.assertGreaterEqual(cooldown_seconds, 90)
+        self.assertEqual(1, removed)
+        self.assertTrue(worker._is_proxy_source_quarantined("D"))
+        self.assertEqual(
+            [{"proxy": good_proxy, "usage_count": 0, "source_tag": "E"}],
+            worker.proxy_pool,
+        )
+
+    def test_prepare_proxy_fetch_plan_respects_pool_shortage_instead_of_overfetching(self):
+        worker = self._build_worker(_FakeDb("noop"))
+
+        plan, meta = worker._prepare_proxy_fetch_plan(
+            [
+                "http://provider.local/get?group=B",
+                "http://provider.local/get?group=C",
+                "http://provider.local/get?group=D",
+                "http://provider.local/get?group=E",
+                "http://provider.local/get?group=F",
+                "http://provider.local/get?group=G",
+            ],
+            current_pool_size=1000,
+        )
+
+        self.assertEqual(1, meta["rounds"])
+        self.assertGreaterEqual(meta["target_total"], 840)
+        self.assertLessEqual(meta["target_total"], 1200)
+        self.assertEqual(6, len(plan))
+
+    def test_prepare_proxy_fetch_plan_caps_active_rounds(self):
+        worker = self._build_worker(_FakeDb("noop"))
+
+        plan, meta = worker._prepare_proxy_fetch_plan(
+            [
+                "http://provider.local/get?group=B",
+                "http://provider.local/get?group=C",
+                "http://provider.local/get?group=D",
+                "http://provider.local/get?group=E",
+                "http://provider.local/get?group=F",
+                "http://provider.local/get?group=G",
+            ],
+            current_pool_size=0,
+        )
+
+        self.assertLessEqual(meta["rounds"], 2)
+        self.assertEqual(meta["rounds"] * 6, len(plan))
+
+    def test_submit_backlog_limit_defaults_to_thread_count_window(self):
+        worker = self._build_worker(_FakeDb("noop"))
+
+        self.assertEqual(2000, worker._resolve_submit_backlog_limit(2000, 0))
+        self.assertEqual(2000, worker._resolve_submit_backlog_limit(2000, 1200))
+
+    def test_dispatch_capacity_can_ramp_to_full_thread_limit(self):
+        worker = self._build_worker(_FakeDb("noop"))
+
+        self.assertEqual(2000, worker._resolve_dispatch_capacity(2000, 0, 2000))
+        self.assertEqual(2000, worker._resolve_dispatch_capacity(2000, 1000, 1000))
+
+    def test_claim_batch_size_defaults_to_thread_count_window(self):
+        worker = self._build_worker(_FakeDb("noop"))
+
+        self.assertEqual(2000, worker._resolve_claim_batch_size(2000, 2000, live_active=0))
+        self.assertEqual(1500, worker._resolve_claim_batch_size(2000, 1500, live_active=800))
+
+    def test_get_proxy_for_step_can_force_single_direct_fallback_once(self):
+        worker = self._build_worker(_FakeDb("noop"))
+        worker._set_current_task_context(domain_id=1, is_step_task=True, task_mode="single_step")
+        worker._set_step_force_direct_once("百度site检测", True)
+
+        proxy = worker._get_proxy_for_step(1, "example.com", "百度site检测")
+
+        self.assertIsNone(proxy)
+        self.assertFalse(worker._consume_step_force_direct_once("百度site检测"))
+
+    def test_retry_external_issue_with_proxy_forces_next_attempt_direct_when_allowed(self):
+        worker = self._build_worker(_FakeDb("noop"))
+        worker._set_current_task_context(domain_id=1, is_step_task=True, task_mode="single_step")
+        worker.remove_proxy = MagicMock()
+        proxy = {"http": "http://bad-proxy", "https": "http://bad-proxy"}
+
+        should_retry = worker._should_retry_external_issue(
+            proxy,
+            "curl: (28) Connection timed out after 552 milliseconds",
+            "百度site检测",
+            "example.com",
+        )
+
+        self.assertTrue(should_retry)
+        worker.remove_proxy.assert_called_once_with(proxy)
+        worker.trigger_proxy_refresh.assert_called_once()
+        self.assertTrue(worker._consume_step_force_direct_once("百度site检测"))
+
+    def test_external_dependency_issue_recognizes_common_chinese_timeout_messages(self):
+        worker = self._build_worker(_FakeDb("noop"))
+
+        self.assertTrue(worker._is_external_dependency_issue("360搜索请求超时"))
+        self.assertTrue(worker._is_external_dependency_issue("百度site检测失败: 连接超时"))
+        self.assertTrue(worker._is_external_dependency_issue("站长之家检测失败: 网络异常"))
+
+    def test_run_detect_aizhan_passes_remaining_budget_to_detector(self):
+        worker = self._build_worker(_FakeDb("noop"))
+        worker._check_proxy_step_retry_budget = MagicMock(return_value=(False, ""))
+        worker._remaining_proxy_step_retry_budget_seconds = MagicMock(return_value=1.7)
+        worker._get_proxy_for_step = MagicMock(return_value=None)
+        worker._mark_detection_degraded = MagicMock()
+        worker._should_blacklist_result = MagicMock(return_value=False)
+        worker._record_step_result = MagicMock()
+        worker._should_rotate_proxy_on_error = MagicMock(return_value=False)
+        worker._is_external_dependency_issue = MagicMock(return_value=False)
+        worker._mark_detection_rejected = MagicMock()
+        worker._mark_detection_failed = MagicMock()
+
+        with patch("detect_worker.aizhan.check_aizhan", return_value=(True, "")) as mock_check:
+            result = worker._run_detect_aizhan(1, "example.com", [])
+
+        self.assertTrue(result)
+        self.assertEqual(1.7, mock_check.call_args.kwargs["budget_seconds"])
+
+
+if __name__ == "__main__":
+    unittest.main()