diff --git a/tests/test_crawler/test_proxy_pool.py b/tests/test_crawler/test_proxy_pool.py new file mode 100644 index 0000000..6700b1b --- /dev/null +++ b/tests/test_crawler/test_proxy_pool.py @@ -0,0 +1,178 @@ +"""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 +""" + +import os +import pytest +from pathlib import Path +from unittest.mock import patch + +# 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