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