diff --git a/kb_workflow/crawler/config.py b/kb_workflow/crawler/config.py new file mode 100644 index 0000000..cdb0630 --- /dev/null +++ b/kb_workflow/crawler/config.py @@ -0,0 +1,141 @@ +"""Crawler configuration + +Handles proxy settings, rate limiting, and retry configuration. +""" + +import os +from urllib.parse import urlparse +from typing import Optional, Dict, Any + +from pydantic import BaseModel, Field + + +class ProxyConfig(BaseModel): + """Proxy configuration for crawler requests. + + Attributes: + server: Proxy server address (e.g., http://127.0.0.1:7890) + username: Optional proxy username for authentication + password: Optional proxy password for authentication + """ + + server: str = Field(..., description="Proxy server address") + username: Optional[str] = Field(None, description="Proxy username") + password: Optional[str] = Field(None, description="Proxy password") + + def to_playwright_format(self) -> Dict[str, Any]: + """Convert to Playwright proxy format. + + Returns: + Dictionary compatible with browser.new_context(proxy=...) + """ + config = {"server": self.server} + if self.username: + config["username"] = self.username + if self.password: + config["password"] = self.password + return config + + +class CrawlerConfig(BaseModel): + """Crawler configuration. + + Handles proxy settings, rate limiting, and retry behavior. + + Attributes: + proxy: Optional proxy configuration + proxy_required: Whether proxy is mandatory (raises error if not configured) + min_delay: Minimum delay between requests (seconds) + max_delay: Maximum delay between requests (seconds) + max_retries: Maximum number of retry attempts + retry_min_wait: Minimum wait time between retries (seconds) + retry_max_wait: Maximum wait time between retries (seconds) + """ + + # Proxy settings + proxy: Optional[ProxyConfig] = Field(None, description="Proxy configuration") + proxy_required: bool = Field(default=True, description="Whether proxy is mandatory") + + # Rate limiting + min_delay: float = Field(default=1.0, description="Minimum delay between requests (seconds)") + max_delay: float = Field(default=3.0, description="Maximum delay between requests (seconds)") + + # Retry settings + max_retries: int = Field(default=3, description="Maximum retry attempts") + retry_min_wait: float = Field(default=2.0, description="Minimum retry wait (seconds)") + retry_max_wait: float = Field(default=10.0, description="Maximum retry wait (seconds)") + + @classmethod + def from_env(cls) -> "CrawlerConfig": + """Load configuration from environment variables. + + Reads HTTP_PROXY or HTTPS_PROXY for proxy settings. + + Returns: + CrawlerConfig instance with proxy settings from environment + """ + proxy = cls._parse_proxy_env() + return cls(proxy=proxy) + + @staticmethod + def _parse_proxy_env() -> Optional[ProxyConfig]: + """Parse proxy configuration from environment variables. + + Checks HTTPS_PROXY first, then HTTP_PROXY. + + Returns: + ProxyConfig if proxy URL found, None otherwise + """ + proxy_url = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY") + + if not proxy_url: + return None + + parsed = urlparse(proxy_url) + server = f"{parsed.scheme}://{parsed.hostname}" + + if parsed.port: + server += f":{parsed.port}" + + return ProxyConfig( + server=server, + username=parsed.username, + password=parsed.password + ) + + def validate_proxy(self) -> None: + """Validate proxy configuration. + + Raises: + ValueError: If proxy_required=True but no proxy is configured + """ + if self.proxy_required and not self.proxy: + raise ValueError( + "Proxy configuration missing! Please set HTTP_PROXY or HTTPS_PROXY environment variable.\n" + "Example: export HTTPS_PROXY=http://127.0.0.1:7890" + ) + + +# Module-level cached config +_crawler_config: Optional[CrawlerConfig] = None + + +def get_crawler_config() -> CrawlerConfig: + """Get crawler configuration singleton. + + Returns: + CrawlerConfig instance loaded from environment variables + """ + global _crawler_config + if _crawler_config is None: + _crawler_config = CrawlerConfig.from_env() + return _crawler_config + + +def reset_crawler_config() -> None: + """Reset the cached crawler configuration. + + Useful when environment variables change during runtime. + """ + global _crawler_config + _crawler_config = None diff --git a/tests/test_crawler/test_proxy_config.py b/tests/test_crawler/test_proxy_config.py new file mode 100644 index 0000000..206538e --- /dev/null +++ b/tests/test_crawler/test_proxy_config.py @@ -0,0 +1,125 @@ +"""Tests for Crawler Proxy Configuration""" + +import os +import pytest +from unittest.mock import patch + +from kb_workflow.crawler.config import ( + CrawlerConfig, + ProxyConfig, + get_crawler_config, + reset_crawler_config, +) + + +class TestProxyConfig: + """Tests for ProxyConfig model.""" + + def test_parse_proxy_from_https_env(self): + """Parse proxy from HTTPS_PROXY environment variable.""" + reset_crawler_config() + with patch.dict(os.environ, {"HTTPS_PROXY": "http://127.0.0.1:7890"}, clear=False): + os.environ.pop("HTTP_PROXY", None) + config = CrawlerConfig.from_env() + assert config.proxy is not None + assert config.proxy.server == "http://127.0.0.1:7890" + + def test_parse_proxy_from_http_env(self): + """Parse proxy from HTTP_PROXY (HTTPS_PROXY takes priority).""" + reset_crawler_config() + with patch.dict(os.environ, { + "HTTP_PROXY": "http://127.0.0.1:8080", + "HTTPS_PROXY": "http://127.0.0.1:7890" + }, clear=False): + config = CrawlerConfig.from_env() + assert config.proxy.server == "http://127.0.0.1:7890" + + def test_parse_proxy_with_auth(self): + """Parse proxy with authentication credentials.""" + reset_crawler_config() + with patch.dict(os.environ, { + "HTTPS_PROXY": "http://user:pass@proxy.example.com:8080" + }, clear=False): + os.environ.pop("HTTP_PROXY", None) + config = CrawlerConfig.from_env() + assert config.proxy.username == "user" + assert config.proxy.password == "pass" + assert config.proxy.server == "http://proxy.example.com:8080" + + def test_no_proxy_configured(self): + """Handle case when no proxy is configured.""" + reset_crawler_config() + # Save current values + old_http = os.environ.get("HTTP_PROXY") + old_https = os.environ.get("HTTPS_PROXY") + + try: + os.environ.pop("HTTP_PROXY", None) + os.environ.pop("HTTPS_PROXY", None) + config = CrawlerConfig.from_env() + assert config.proxy is None + finally: + # Restore + if old_http: + os.environ["HTTP_PROXY"] = old_http + if old_https: + os.environ["HTTPS_PROXY"] = old_https + + def test_proxy_required_validation_fails(self): + """Raise ValueError when proxy required but not configured.""" + config = CrawlerConfig(proxy=None, proxy_required=True) + with pytest.raises(ValueError, match="Proxy configuration missing"): + config.validate_proxy() + + def test_proxy_required_validation_passes(self): + """Pass validation when proxy is configured and required.""" + config = CrawlerConfig( + proxy=ProxyConfig(server="http://127.0.0.1:7890"), + proxy_required=True + ) + # Should not raise exception + config.validate_proxy() + + def test_proxy_not_required(self): + """Pass validation when proxy is not required.""" + config = CrawlerConfig(proxy=None, proxy_required=False) + # Should not raise exception + config.validate_proxy() + + def test_to_playwright_format(self): + """Convert to Playwright proxy format.""" + proxy = ProxyConfig( + server="http://127.0.0.1:7890", + username="user", + password="pass" + ) + result = proxy.to_playwright_format() + assert result == { + "server": "http://127.0.0.1:7890", + "username": "user", + "password": "pass" + } + + def test_to_playwright_format_no_auth(self): + """Convert to Playwright format without authentication.""" + proxy = ProxyConfig(server="http://127.0.0.1:7890") + result = proxy.to_playwright_format() + assert result == {"server": "http://127.0.0.1:7890"} + + def test_get_crawler_config_singleton(self): + """Test that get_crawler_config returns cached instance.""" + reset_crawler_config() + with patch.dict(os.environ, {"HTTPS_PROXY": "http://127.0.0.1:7890"}, clear=False): + os.environ.pop("HTTP_PROXY", None) + config1 = get_crawler_config() + config2 = get_crawler_config() + assert config1 is config2 + + def test_default_values(self): + """Test default configuration values.""" + config = CrawlerConfig(proxy=None, proxy_required=False) + assert config.min_delay == 1.0 + assert config.max_delay == 3.0 + assert config.max_retries == 3 + assert config.retry_min_wait == 2.0 + assert config.retry_max_wait == 10.0