226 lines
9.5 KiB
Python
226 lines
9.5 KiB
Python
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()
|