diff --git a/kb_workflow/crawler/zhihu_crawler.py b/kb_workflow/crawler/zhihu_crawler.py index e343f5f..77b029b 100644 --- a/kb_workflow/crawler/zhihu_crawler.py +++ b/kb_workflow/crawler/zhihu_crawler.py @@ -8,12 +8,14 @@ Supports: - Random request delay - Retry with exponential backoff - Automatic x-zse-96 signature handling via browser +- ProxyPool integration for multi-proxy rotation """ import random import re import time from typing import Optional, Dict, Any +from urllib.parse import urlparse from playwright.sync_api import Route, Response, sync_playwright from tenacity import Retrying, stop_after_attempt, wait_exponential, RetryError @@ -32,6 +34,7 @@ from .models import ( CrawlResult, ) from .config import CrawlerConfig, get_crawler_config +from .proxy_pool import ProxyPool, ProxyEntry class ZhihuApiInterceptor: @@ -104,7 +107,8 @@ class ZhihuCrawler: self, headless: bool = True, timeout: int = 30, - config: Optional[CrawlerConfig] = None + config: Optional[CrawlerConfig] = None, + proxy_pool: Optional[ProxyPool] = None ): """Initialize ZhihuCrawler. @@ -112,10 +116,13 @@ class ZhihuCrawler: headless: Whether to run browser in headless mode timeout: Request timeout in seconds config: Optional CrawlerConfig instance (uses environment if not provided) + proxy_pool: Optional ProxyPool instance for multi-proxy rotation """ self.headless = headless self.timeout = timeout self.config = config or get_crawler_config() + self.proxy_pool = proxy_pool + self._current_proxy: Optional[ProxyEntry] = None self.browser: Optional[PlaywrightBrowser] = None self.interceptor: Optional[ZhihuApiInterceptor] = None # Internal browser management for proxy support @@ -148,6 +155,7 @@ class ZhihuCrawler: """Start browser with proxy configuration. Validates proxy if required and applies to browser context. + Prioritizes proxy_pool over config.proxy for multi-proxy support. Args: cookie: Cookie string for authentication @@ -166,8 +174,15 @@ class ZhihuCrawler: 'user_agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36' } - # Apply proxy if configured - if self.config.proxy: + # Prioritize proxy pool over single proxy config + if self.proxy_pool and self.proxy_pool.has_proxies(): + proxy_entry = self.proxy_pool.get_next_proxy() + if proxy_entry: + proxy_config = self._parse_proxy_url(proxy_entry.url) + context_options["proxy"] = proxy_config + self._current_proxy = proxy_entry + elif self.config.proxy: + # Backward compatibility: use single proxy configuration context_options["proxy"] = self.config.proxy.to_playwright_format() # Block Service Worker to prevent interference with API interception @@ -182,6 +197,40 @@ class ZhihuCrawler: self._page = self._context.new_page() self._page.set_default_timeout(self.timeout * 1000) + @staticmethod + def _parse_proxy_url(url: str) -> Dict[str, Any]: + """Parse proxy URL into Playwright format. + + Args: + url: Proxy URL (e.g., http://user:pass@proxy.example.com:8080) + + Returns: + Dictionary with server, username, password for Playwright + """ + parsed = urlparse(url) + server = f"{parsed.scheme}://{parsed.hostname}" + if parsed.port: + server += f":{parsed.port}" + config: Dict[str, Any] = {"server": server} + if parsed.username: + config["username"] = parsed.username + if parsed.password: + config["password"] = parsed.password + return config + + def _on_request_success(self) -> None: + """Notify proxy pool of successful request.""" + if self._current_proxy: + self._current_proxy.record_success() + + def _on_request_failure(self) -> None: + """Notify proxy pool of failed request.""" + if self._current_proxy: + self._current_proxy.record_failure( + max_fail_count=self.proxy_pool.config.max_fail_count + if self.proxy_pool else 3 + ) + def _add_cookies(self, cookie_str: str) -> None: """Add cookies to browser context. diff --git a/tests/test_crawler/test_proxy_pool.py b/tests/test_crawler/test_proxy_pool.py index f6f5262..2ffeca7 100644 --- a/tests/test_crawler/test_proxy_pool.py +++ b/tests/test_crawler/test_proxy_pool.py @@ -6,12 +6,14 @@ Tests cover: - 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 +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 @@ -304,3 +306,70 @@ class TestProxyPoolHealthCheck: 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