kb-skills/tests/test_crawler/test_rate_limiter.py

"""Tests for Rate Limiter and Retry Logic"""

import pytest
import time
from unittest.mock import Mock, patch, MagicMock

from kb_workflow.crawler.zhihu_crawler import ZhihuCrawler
from kb_workflow.crawler.config import CrawlerConfig


class TestRateLimiter:
    """Tests for random delay and retry logic."""

    def test_random_delay_applied(self):
        """Request has random delay between min_delay and max_delay."""
        config = CrawlerConfig(
            min_delay=0.1,
            max_delay=0.2,
            proxy_required=False
        )
        crawler = ZhihuCrawler(config=config)

        start = time.time()
        crawler._apply_random_delay()
        elapsed = time.time() - start

        # Delay should be within range (with small margin for timing)
        assert 0.1 <= elapsed <= 0.3

    def test_delay_is_random(self):
        """Delay is randomized, not fixed."""
        config = CrawlerConfig(
            min_delay=0.05,
            max_delay=0.15,
            proxy_required=False
        )
        crawler = ZhihuCrawler(config=config)

        delays = []
        for _ in range(10):
            start = time.time()
            crawler._apply_random_delay()
            delays.append(time.time() - start)

        # Delays should have variation
        assert max(delays) - min(delays) > 0.01

    def test_exponential_backoff_retry(self):
        """Retry uses exponential backoff with increasing intervals."""
        config = CrawlerConfig(
            max_retries=3,
            retry_min_wait=0.1,
            retry_max_wait=1.0,
            proxy_required=False
        )
        crawler = ZhihuCrawler(config=config)

        call_times = []

        def failing_func():
            call_times.append(time.time())
            # Simulate a network error - raise exception to trigger retry
            raise ConnectionError("Simulated network error")

        # Mock the delay to speed up test
        with patch.object(crawler, '_apply_random_delay'):
            result = crawler._crawl_with_retry(failing_func)

        # Should have failed after retries
        assert result.success is False
        # Should have attempted max_retries times
        assert len(call_times) == 3  # Initial + 2 retries

    def test_retry_success_after_failure(self):
        """Retry succeeds if function eventually succeeds."""
        config = CrawlerConfig(
            max_retries=3,
            retry_min_wait=0.05,
            retry_max_wait=0.2,
            proxy_required=False
        )
        crawler = ZhihuCrawler(config=config)

        call_count = [0]

        def eventually_succeeds():
            call_count[0] += 1
            if call_count[0] < 2:
                raise Exception("Simulated failure")
            return Mock(success=True, content_type="test")

        # Mock _apply_random_delay to avoid actual delay
        with patch.object(crawler, '_apply_random_delay'):
            result = crawler._crawl_with_retry(eventually_succeeds)

        assert result.success is True
        assert call_count[0] == 2  # Failed once, succeeded on second try

    def test_crawl_with_retry_returns_result(self):
        """_crawl_with_retry returns CrawlResult from wrapped function."""
        config = CrawlerConfig(proxy_required=False)
        crawler = ZhihuCrawler(config=config)

        from kb_workflow.crawler.models import CrawlResult

        expected_result = CrawlResult(success=True, content_type="test")

        def success_func():
            return expected_result

        with patch.object(crawler, '_apply_random_delay'):
            result = crawler._crawl_with_retry(success_func)

        assert result.success is True
        assert result.content_type == "test"


class TestProxyIntegration:
    """Tests for proxy configuration integration."""

    def test_no_proxy_with_proxy_required_raises(self):
        """Raises ValueError when proxy_required=True but no proxy configured."""
        config = CrawlerConfig(proxy=None, proxy_required=True)
        crawler = ZhihuCrawler(config=config)

        with pytest.raises(ValueError, match="Proxy configuration missing"):
            crawler._start_browser_with_proxy("test_cookie")

    def test_no_proxy_with_proxy_not_required_passes(self):
        """Succeeds when proxy_required=False and no proxy configured."""
        config = CrawlerConfig(proxy=None, proxy_required=False)
        crawler = ZhihuCrawler(config=config)

        # Mock the browser startup to avoid actual browser launch
        with patch.object(crawler, '_playwright') as mock_pw:
            mock_pb_instance = MagicMock()
            mock_pw.return_value = mock_pb_instance
            mock_pb_instance.chromium.launch.return_value = MagicMock()

            # Should not raise
            crawler._start_browser_with_proxy("test_cookie")

    def test_proxy_applied_to_browser_context(self):
        """Proxy is applied to browser context when configured."""
        from kb_workflow.crawler.config import ProxyConfig

        config = CrawlerConfig(
            proxy=ProxyConfig(server="http://127.0.0.1:7890"),
            proxy_required=True
        )
        crawler = ZhihuCrawler(config=config)

        # Mock browser components
        with patch('kb_workflow.crawler.zhihu_crawler.sync_playwright') as mock_sp:
            mock_pw = MagicMock()
            mock_sp.return_value.start.return_value = mock_pw

            mock_browser = MagicMock()
            mock_pw.chromium.launch.return_value = mock_browser

            mock_context = MagicMock()
            mock_browser.new_context.return_value = mock_context

            crawler._start_browser_with_proxy("test_cookie")

            # Verify proxy was passed to new_context
            call_kwargs = mock_browser.new_context.call_args[1]
            assert "proxy" in call_kwargs
            assert call_kwargs["proxy"]["server"] == "http://127.0.0.1:7890"