websurfer/test_tor_fallback.py

import asyncio
import unittest
from unittest.mock import MagicMock, AsyncMock, patch

from crawler.config import CrawlerConfig
from crawler.engine import Crawler


class TestTorTimeoutFallback(unittest.IsolatedAsyncioTestCase):

    def test_config_defaults_and_env(self):
        config = CrawlerConfig(start_url="https://example.com")
        self.assertEqual(config.tor_max_consecutive_timeouts, 5)

        dict_data = config.to_dict()
        self.assertEqual(dict_data["tor_max_consecutive_timeouts"], 5)

        dict_data["tor_max_consecutive_timeouts"] = 3
        restored = CrawlerConfig.from_dict(dict_data)
        self.assertEqual(restored.tor_max_consecutive_timeouts, 3)

    async def test_consecutive_timeout_fallback_trigger(self):
        config = CrawlerConfig(
            start_url="https://example.com",
            use_tor=True,
            tor_max_consecutive_timeouts=5,
        )
        events = []
        async def on_event(ev):
            events.append(ev)

        crawler = Crawler(config=config, event_callback=on_event)
        # Mock tor_manager so we don't need real Tor daemon for unit testing
        crawler.tor_manager = MagicMock()
        crawler.tor_manager.renew_circuit = MagicMock(return_value=True)

        mock_context = MagicMock()

        # 1. First 4 consecutive ERR_TIMED_OUT errors
        for i in range(1, 5):
            await crawler._on_page_result(
                context=mock_context,
                url=f"https://example.com/page{i}",
                source_url="https://example.com",
                success=False,
                error_msg="page.goto: net::ERR_TIMED_OUT at https://example.com",
            )
            self.assertEqual(crawler._consecutive_timeouts, i)
            self.assertEqual(crawler.tor_manager.renew_circuit.call_count, 0)

        # 2. 5th ERR_TIMED_OUT error triggers fallback circuit rotation!
        await crawler._on_page_result(
            context=mock_context,
            url="https://example.com/page5",
            source_url="https://example.com",
            success=False,
            error_msg="page.goto: net::ERR_TIMED_OUT at https://example.com",
        )
        # Counter should be reset to 0 after rotation
        self.assertEqual(crawler._consecutive_timeouts, 0)
        self.assertEqual(crawler.tor_manager.renew_circuit.call_count, 1)

    async def test_success_resets_counter(self):
        config = CrawlerConfig(
            start_url="https://example.com",
            use_tor=True,
            tor_max_consecutive_timeouts=5,
        )
        crawler = Crawler(config=config)
        crawler.tor_manager = MagicMock()
        crawler.tor_manager.renew_circuit = MagicMock(return_value=True)

        mock_context = MagicMock()

        # 3 timeouts
        for i in range(1, 4):
            await crawler._on_page_result(
                context=mock_context,
                url=f"https://example.com/page{i}",
                source_url=None,
                success=False,
                error_msg="net::ERR_TIMED_OUT",
            )
        self.assertEqual(crawler._consecutive_timeouts, 3)

        # 1 success -> should reset counter to 0
        await crawler._on_page_result(
            context=mock_context,
            url="https://example.com/success_page",
            source_url=None,
            success=True,
            error_msg=None,
        )
        self.assertEqual(crawler._consecutive_timeouts, 0)
        self.assertEqual(crawler.tor_manager.renew_circuit.call_count, 0)

    async def test_non_timeout_error_resets_counter(self):
        config = CrawlerConfig(
            start_url="https://example.com",
            use_tor=True,
            tor_max_consecutive_timeouts=5,
        )
        crawler = Crawler(config=config)
        crawler.tor_manager = MagicMock()

        mock_context = MagicMock()

        # 2 timeouts
        for i in range(2):
            await crawler._on_page_result(
                context=mock_context,
                url=f"https://example.com/page{i}",
                source_url=None,
                success=False,
                error_msg="net::ERR_TIMED_OUT",
            )
        self.assertEqual(crawler._consecutive_timeouts, 2)

        # HTTP 404 error -> non-timeout/non-rate-limit resets counter
        await crawler._on_page_result(
            context=mock_context,
            url="https://example.com/missing",
            source_url=None,
            success=False,
            error_msg="HTTP 404",
        )
        self.assertEqual(crawler._consecutive_timeouts, 0)

    def test_user_agent_config(self):
        config_random = CrawlerConfig(start_url="https://example.com", user_agent="random")
        ua1 = config_random.get_user_agent()
        self.assertIsNotNone(ua1)
        self.assertIn("Mozilla/5.0", ua1)

        config_custom = CrawlerConfig(start_url="https://example.com", user_agent="MyCustomBot/1.0")
        self.assertEqual(config_custom.get_user_agent(), "MyCustomBot/1.0")

        config_default = CrawlerConfig(start_url="https://example.com", user_agent="default")
        self.assertIsNone(config_default.get_user_agent())

    async def test_consecutive_rate_limit_fallback_trigger(self):
        config = CrawlerConfig(
            start_url="https://example.com",
            use_tor=True,
            tor_max_consecutive_timeouts=3,
        )
        crawler = Crawler(config=config)
        crawler.tor_manager = MagicMock()
        crawler.tor_manager.renew_circuit = MagicMock(return_value=True)

        mock_context = MagicMock()

        # 2 HTTP 503 errors
        for i in range(2):
            await crawler._on_page_result(
                context=mock_context,
                url=f"https://example.com/page{i}",
                source_url="https://example.com",
                success=False,
                error_msg="HTTP 503 (Service Temporarily Unavailable / Rate Limited) [Retry-After: 154s]",
            )
        self.assertEqual(crawler._consecutive_timeouts, 2)
        self.assertEqual(crawler.tor_manager.renew_circuit.call_count, 0)

        # 3rd HTTP 503 error triggers circuit renewal
        await crawler._on_page_result(
            context=mock_context,
            url="https://example.com/page2",
            source_url="https://example.com",
            success=False,
            error_msg="HTTP 503 (Service Temporarily Unavailable / Rate Limited) [Retry-After: 154s]",
        )
        self.assertEqual(crawler._consecutive_timeouts, 0)
        self.assertEqual(crawler.tor_manager.renew_circuit.call_count, 1)


if __name__ == "__main__":
    unittest.main()