starmen1265
b66bcc08b1
- Add proxy_pool parameter to ZhihuCrawler constructor - Add _current_proxy tracking for health check feedback - Implement _parse_proxy_url static method for URL parsing - Implement _on_request_success/failure for proxy health feedback - Prioritize proxy_pool over config.proxy for multi-proxy support - Maintain backward compatibility with single proxy config Tests: - test_crawler_accepts_proxy_pool - test_crawler_without_proxy_pool (backward compatibility) - test_parse_proxy_url (HTTP, auth, SOCKS5) - test_proxy_pool_takes_priority_over_config Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
376 lines
12 KiB
Python
376 lines
12 KiB
Python
"""Tests for proxy pool functionality.
|
|
|
|
Tests cover:
|
|
- Loading proxy pool from YAML file
|
|
- Loading proxy from environment variable (backward compatibility)
|
|
- Round-robin proxy rotation
|
|
- Empty pool handling
|
|
- Single proxy handling
|
|
- Health check and cooldown recovery
|
|
- ZhihuCrawler integration with ProxyPool
|
|
"""
|
|
|
|
import os
|
|
import pytest
|
|
from pathlib import Path
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
# Import will fail until proxy_pool module is implemented (RED phase)
|
|
from kb_workflow.crawler.proxy_pool import ProxyPool, ProxyEntry, ProxyPoolConfig
|
|
|
|
|
|
class TestLoadFromYaml:
|
|
"""Tests for loading proxy pool from YAML file."""
|
|
|
|
def test_load_from_yaml(self, tmp_path: Path):
|
|
"""Test loading proxy pool from YAML file with multiple proxies."""
|
|
# Create temporary YAML file
|
|
yaml_content = """
|
|
proxies:
|
|
- "http://127.0.0.1:7890"
|
|
- "http://user:pass@proxy.example.com:8080"
|
|
- "socks5://127.0.0.1:1080"
|
|
cooldown_seconds: 300
|
|
max_fail_count: 3
|
|
"""
|
|
yaml_file = tmp_path / ".kb_proxies.yaml"
|
|
yaml_file.write_text(yaml_content)
|
|
|
|
# Load proxy pool
|
|
pool = ProxyPool.from_yaml(yaml_file)
|
|
|
|
# Verify
|
|
assert pool.has_proxies()
|
|
assert pool.config.cooldown_seconds == 300
|
|
assert pool.config.max_fail_count == 3
|
|
|
|
def test_load_from_yaml_not_found(self, tmp_path: Path):
|
|
"""Test loading from non-existent YAML file returns empty pool."""
|
|
non_existent = tmp_path / "non_existent.yaml"
|
|
pool = ProxyPool.from_yaml(non_existent)
|
|
|
|
# Should return empty pool, not raise exception
|
|
assert not pool.has_proxies()
|
|
|
|
def test_load_from_yaml_empty_file(self, tmp_path: Path):
|
|
"""Test loading from empty YAML file returns empty pool."""
|
|
empty_file = tmp_path / "empty.yaml"
|
|
empty_file.write_text("")
|
|
|
|
pool = ProxyPool.from_yaml(empty_file)
|
|
assert not pool.has_proxies()
|
|
|
|
|
|
class TestLoadFromEnv:
|
|
"""Tests for loading proxy from environment variable."""
|
|
|
|
def test_load_from_env_https_proxy(self):
|
|
"""Test loading proxy from HTTPS_PROXY environment variable."""
|
|
with patch.dict(os.environ, {"HTTPS_PROXY": "http://127.0.0.1:7890"}, clear=False):
|
|
pool = ProxyPool.from_env_fallback()
|
|
assert pool.has_proxies()
|
|
|
|
def test_load_from_env_http_proxy(self):
|
|
"""Test loading proxy from HTTP_PROXY environment variable (fallback)."""
|
|
# Clear HTTPS_PROXY if set
|
|
env = os.environ.copy()
|
|
env.pop("HTTPS_PROXY", None)
|
|
env["HTTP_PROXY"] = "http://127.0.0.1:8080"
|
|
|
|
with patch.dict(os.environ, env, clear=True):
|
|
pool = ProxyPool.from_env_fallback()
|
|
assert pool.has_proxies()
|
|
|
|
def test_load_from_env_no_proxy(self):
|
|
"""Test loading without any proxy environment variable."""
|
|
env = os.environ.copy()
|
|
env.pop("HTTPS_PROXY", None)
|
|
env.pop("HTTP_PROXY", None)
|
|
|
|
with patch.dict(os.environ, env, clear=True):
|
|
pool = ProxyPool.from_env_fallback()
|
|
assert not pool.has_proxies()
|
|
|
|
|
|
class TestRoundRobin:
|
|
"""Tests for round-robin proxy rotation."""
|
|
|
|
def test_round_robin(self):
|
|
"""Test round-robin rotation with multiple proxies."""
|
|
config = ProxyPoolConfig(
|
|
proxies=[
|
|
"http://proxy-a:8080",
|
|
"http://proxy-b:8080",
|
|
"http://proxy-c:8080",
|
|
]
|
|
)
|
|
pool = ProxyPool(config)
|
|
|
|
# Call get_next_proxy 6 times, should cycle A-B-C-A-B-C
|
|
results = []
|
|
for _ in range(6):
|
|
proxy = pool.get_next_proxy()
|
|
if proxy:
|
|
results.append(proxy.url)
|
|
|
|
expected = [
|
|
"http://proxy-a:8080",
|
|
"http://proxy-b:8080",
|
|
"http://proxy-c:8080",
|
|
"http://proxy-a:8080",
|
|
"http://proxy-b:8080",
|
|
"http://proxy-c:8080",
|
|
]
|
|
assert results == expected
|
|
|
|
def test_round_robin_consistency(self):
|
|
"""Test that round-robin is consistent across multiple cycles."""
|
|
config = ProxyPoolConfig(
|
|
proxies=[
|
|
"http://proxy-1:8080",
|
|
"http://proxy-2:8080",
|
|
]
|
|
)
|
|
pool = ProxyPool(config)
|
|
|
|
# First cycle
|
|
first_cycle = [pool.get_next_proxy().url for _ in range(2)]
|
|
# Second cycle
|
|
second_cycle = [pool.get_next_proxy().url for _ in range(2)]
|
|
|
|
assert first_cycle == second_cycle
|
|
|
|
|
|
class TestEmptyPool:
|
|
"""Tests for empty proxy pool handling."""
|
|
|
|
def test_empty_pool(self):
|
|
"""Test that get_next_proxy returns None for empty pool."""
|
|
config = ProxyPoolConfig(proxies=[])
|
|
pool = ProxyPool(config)
|
|
|
|
assert pool.get_next_proxy() is None
|
|
assert not pool.has_proxies()
|
|
|
|
|
|
class TestSingleProxy:
|
|
"""Tests for single proxy handling."""
|
|
|
|
def test_single_proxy(self):
|
|
"""Test that single proxy is always returned."""
|
|
config = ProxyPoolConfig(proxies=["http://single-proxy:8080"])
|
|
pool = ProxyPool(config)
|
|
|
|
# Should always return the same proxy
|
|
for _ in range(5):
|
|
proxy = pool.get_next_proxy()
|
|
assert proxy.url == "http://single-proxy:8080"
|
|
|
|
|
|
class TestProxyEntry:
|
|
"""Tests for ProxyEntry class."""
|
|
|
|
def test_proxy_entry_defaults(self):
|
|
"""Test ProxyEntry default values."""
|
|
entry = ProxyEntry(url="http://proxy:8080")
|
|
|
|
assert entry.url == "http://proxy:8080"
|
|
assert entry.fail_count == 0
|
|
assert entry.last_fail_time is None
|
|
assert entry.is_available is True
|
|
|
|
def test_fail_threshold(self):
|
|
"""Test that proxy becomes unavailable after 3 consecutive failures."""
|
|
entry = ProxyEntry(url="http://proxy:8080")
|
|
|
|
# First two failures should not mark as unavailable
|
|
entry.record_failure(max_fail_count=3)
|
|
assert entry.fail_count == 1
|
|
assert entry.is_available is True
|
|
|
|
entry.record_failure(max_fail_count=3)
|
|
assert entry.fail_count == 2
|
|
assert entry.is_available is True
|
|
|
|
# Third failure should mark as unavailable
|
|
entry.record_failure(max_fail_count=3)
|
|
assert entry.fail_count == 3
|
|
assert entry.is_available is False
|
|
|
|
def test_success_resets_count(self):
|
|
"""Test that recording success resets failure count and restores availability."""
|
|
entry = ProxyEntry(url="http://proxy:8080")
|
|
|
|
# Record some failures
|
|
entry.record_failure(max_fail_count=3)
|
|
entry.record_failure(max_fail_count=3)
|
|
assert entry.fail_count == 2
|
|
|
|
# Record success
|
|
entry.record_success()
|
|
|
|
assert entry.fail_count == 0
|
|
assert entry.is_available is True
|
|
|
|
def test_cooldown_recovery(self):
|
|
"""Test that failed proxy recovers after cooldown period."""
|
|
import time
|
|
|
|
entry = ProxyEntry(url="http://proxy:8080")
|
|
|
|
# Mark as unavailable
|
|
entry.record_failure(max_fail_count=3)
|
|
entry.record_failure(max_fail_count=3)
|
|
entry.record_failure(max_fail_count=3)
|
|
assert entry.is_available is False
|
|
|
|
# Wait for cooldown (using short cooldown for testing)
|
|
time.sleep(1)
|
|
|
|
# Check recovery with 1 second cooldown
|
|
result = entry.check_recovery(cooldown_seconds=1)
|
|
assert result is True
|
|
assert entry.is_available is True
|
|
assert entry.fail_count == 0
|
|
|
|
def test_cooldown_not_elapsed(self):
|
|
"""Test that failed proxy does not recover before cooldown period."""
|
|
entry = ProxyEntry(url="http://proxy:8080")
|
|
|
|
# Mark as unavailable
|
|
entry.record_failure(max_fail_count=3)
|
|
entry.record_failure(max_fail_count=3)
|
|
entry.record_failure(max_fail_count=3)
|
|
assert entry.is_available is False
|
|
|
|
# Check recovery immediately with 60 second cooldown
|
|
result = entry.check_recovery(cooldown_seconds=60)
|
|
assert result is False
|
|
assert entry.is_available is False
|
|
|
|
|
|
class TestProxyPoolHealthCheck:
|
|
"""Tests for ProxyPool health check and availability methods."""
|
|
|
|
def test_all_unavailable_fallback(self):
|
|
"""Test that get_next_proxy returns first proxy when all are unavailable."""
|
|
config = ProxyPoolConfig(
|
|
proxies=[
|
|
"http://proxy-a:8080",
|
|
"http://proxy-b:8080",
|
|
"http://proxy-c:8080",
|
|
],
|
|
max_fail_count=3
|
|
)
|
|
pool = ProxyPool(config)
|
|
|
|
# Mark all proxies as unavailable
|
|
for _ in range(3):
|
|
pool._entries[0].record_failure(max_fail_count=3)
|
|
pool._entries[1].record_failure(max_fail_count=3)
|
|
pool._entries[2].record_failure(max_fail_count=3)
|
|
|
|
# All should be unavailable
|
|
assert pool._entries[0].is_available is False
|
|
assert pool._entries[1].is_available is False
|
|
assert pool._entries[2].is_available is False
|
|
|
|
# Should fallback to first proxy
|
|
proxy = pool.get_next_proxy()
|
|
assert proxy is not None
|
|
assert proxy.url == "http://proxy-a:8080"
|
|
|
|
def test_get_available_count(self):
|
|
"""Test that get_available_count returns correct number of available proxies."""
|
|
config = ProxyPoolConfig(
|
|
proxies=[
|
|
"http://proxy-a:8080",
|
|
"http://proxy-b:8080",
|
|
"http://proxy-c:8080",
|
|
],
|
|
max_fail_count=3
|
|
)
|
|
pool = ProxyPool(config)
|
|
|
|
# Initially all available
|
|
assert pool.get_available_count() == 3
|
|
|
|
# Mark one as unavailable
|
|
for _ in range(3):
|
|
pool._entries[0].record_failure(max_fail_count=3)
|
|
|
|
assert pool.get_available_count() == 2
|
|
|
|
# Mark another as unavailable
|
|
for _ in range(3):
|
|
pool._entries[1].record_failure(max_fail_count=3)
|
|
|
|
assert pool.get_available_count() == 1
|
|
|
|
|
|
class TestZhihuCrawlerProxyPoolIntegration:
|
|
"""Tests for ZhihuCrawler integration with ProxyPool."""
|
|
|
|
def test_crawler_accepts_proxy_pool(self):
|
|
"""Test that ZhihuCrawler accepts proxy_pool parameter."""
|
|
from kb_workflow.crawler.zhihu_crawler import ZhihuCrawler
|
|
|
|
config = ProxyPoolConfig(
|
|
proxies=["http://proxy-a:8080", "http://proxy-b:8080"]
|
|
)
|
|
pool = ProxyPool(config)
|
|
|
|
crawler = ZhihuCrawler(proxy_pool=pool)
|
|
|
|
assert crawler.proxy_pool is pool
|
|
assert crawler.proxy_pool.has_proxies()
|
|
|
|
def test_crawler_without_proxy_pool(self):
|
|
"""Test that ZhihuCrawler works without proxy_pool (backward compatibility)."""
|
|
from kb_workflow.crawler.zhihu_crawler import ZhihuCrawler
|
|
|
|
crawler = ZhihuCrawler()
|
|
|
|
assert crawler.proxy_pool is None
|
|
|
|
def test_parse_proxy_url(self):
|
|
"""Test _parse_proxy_url static method."""
|
|
from kb_workflow.crawler.zhihu_crawler import ZhihuCrawler
|
|
|
|
# Test HTTP proxy without auth
|
|
result = ZhihuCrawler._parse_proxy_url("http://127.0.0.1:7890")
|
|
assert result == {"server": "http://127.0.0.1:7890"}
|
|
|
|
# Test HTTP proxy with auth
|
|
result = ZhihuCrawler._parse_proxy_url("http://user:pass@proxy.example.com:8080")
|
|
assert result == {
|
|
"server": "http://proxy.example.com:8080",
|
|
"username": "user",
|
|
"password": "pass"
|
|
}
|
|
|
|
# Test SOCKS5 proxy
|
|
result = ZhihuCrawler._parse_proxy_url("socks5://127.0.0.1:1080")
|
|
assert result == {"server": "socks5://127.0.0.1:1080"}
|
|
|
|
def test_proxy_pool_takes_priority_over_config(self):
|
|
"""Test that proxy_pool takes priority over config.proxy."""
|
|
from kb_workflow.crawler.zhihu_crawler import ZhihuCrawler
|
|
from kb_workflow.crawler.config import CrawlerConfig, ProxyConfig
|
|
|
|
# Create config with single proxy
|
|
single_proxy_config = CrawlerConfig(
|
|
proxy=ProxyConfig(server="http://config-proxy:9999")
|
|
)
|
|
|
|
# Create proxy pool with different proxies
|
|
pool = ProxyPool(ProxyPoolConfig(
|
|
proxies=["http://pool-proxy:8080"]
|
|
))
|
|
|
|
crawler = ZhihuCrawler(config=single_proxy_config, proxy_pool=pool)
|
|
|
|
# Verify proxy_pool is set
|
|
assert crawler.proxy_pool is pool
|
|
assert crawler.config.proxy is not None
|