learningcircuit--local-deep-research/tests/rate_limiting/test_rate_limiting.py

"""
Tests for the adaptive rate limiting system.
"""

import os
import unittest
from unittest.mock import patch

import pytest

from local_deep_research.web_search_engines.rate_limiting import (
    AdaptiveRateLimitTracker,
    RateLimitError,
)


class TestAdaptiveRateLimitTracker(unittest.TestCase):
    """Test the AdaptiveRateLimitTracker class."""

    @classmethod
    def setUpClass(cls):
        """Set up class-level fixtures."""
        # Tables are created automatically via SQLAlchemy models
        pass

    def setUp(self):
        """Set up test fixtures."""
        # Note: Using the main database - in a real test environment
        # you'd want to mock the database session

        # Force settings to use default values
        # The deprecated function has been removed
        # Create settings snapshot with rate limiting enabled for tests
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"}
        }

        # Use programmatic_mode=True for tests to avoid needing user context
        self.tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # Skip database cleanup in CI to avoid timeouts
        if os.environ.get("CI") != "true":
            # Clean up any existing test data before each test
            test_engines = [
                "TestEngine",
                "TestEngine_GetStats",
                "TestEngine_Reset",
                "SearXNGSearchEngine",
            ]
            for engine in test_engines:
                try:
                    self.tracker.reset_engine(engine)
                except Exception:
                    pass

    def tearDown(self):
        """Clean up test fixtures."""
        # Skip database cleanup in CI to avoid timeouts
        if os.environ.get("CI") != "true":
            # Clean up any test data
            test_engines = [
                "TestEngine",
                "TestEngine_GetStats",
                "TestEngine_Reset",
                "SearXNGSearchEngine",
            ]
            for engine in test_engines:
                try:
                    self.tracker.reset_engine(engine)
                except Exception:
                    pass

    @pytest.mark.timeout(30)
    def test_get_wait_time_new_engine(self):
        """Test getting wait time for a new engine.

        Uses timeout marker instead of skipping to fail fast if DB operations hang.
        """
        # Reset any existing data for test engines
        try:
            self.tracker.reset_engine("TestEngine")
            self.tracker.reset_engine("SearXNGSearchEngine")
        except Exception:
            pass

        # Test default engine (unknown)
        wait_time = self.tracker.get_wait_time("TestEngine")
        self.assertEqual(wait_time, 0.1)  # Default optimistic

        # Test SearXNG (self-hosted default)
        # Clear from current estimates to force default
        if "SearXNGSearchEngine" in self.tracker.current_estimates:
            del self.tracker.current_estimates["SearXNGSearchEngine"]
        searxng_wait = self.tracker.get_wait_time("SearXNGSearchEngine")
        self.assertEqual(searxng_wait, 0.1)  # Very optimistic for self-hosted

    @pytest.mark.timeout(30)
    def test_record_outcome_and_learning(self):
        """Test recording outcomes and learning from them.

        Uses timeout marker instead of skipping to fail fast if DB operations hang.
        """
        engine_type = "TestEngine"

        # Record several successful attempts with different wait times
        successful_waits = [2.0, 2.5, 3.0, 2.2, 2.8]
        for i, wait_time in enumerate(successful_waits):
            self.tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=True,
                retry_count=1,
            )

        # The tracker should have learned from these attempts
        self.assertIn(engine_type, self.tracker.current_estimates)

        # Get new wait time - should be influenced by successful attempts
        new_wait_time = self.tracker.get_wait_time(engine_type)
        self.assertGreater(new_wait_time, 0)

    def test_record_failure_increases_wait_time(self):
        """Test that failures increase the wait time when all attempts fail."""
        engine_type = "TestEngine"

        # Record some initial successful attempts (need at least 3 for estimate creation)
        for wait_time in [2.0, 2.5, 3.0]:
            self.tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=True,
                retry_count=1,
            )

        initial_estimate = self.tracker.current_estimates[engine_type]["base"]

        # Reset and record only failures to test failure handling
        self.tracker.recent_attempts[engine_type] = (
            self.tracker.recent_attempts[engine_type].__class__(
                maxlen=self.tracker.memory_window
            )
        )

        # Record only failures (this should increase wait time)
        for wait_time in [2.0, 2.5, 3.0]:
            self.tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=False,
                retry_count=2,
                error_type="RateLimitError",
            )

        # Base wait time should increase after all failures
        new_estimate = self.tracker.current_estimates[engine_type]["base"]
        self.assertGreater(new_estimate, initial_estimate)

    @pytest.mark.skip(
        reason="Skip database persistence test - requires user context which is not available in test environment",
    )
    def test_persistence(self):
        """Test that estimates are persisted across instances."""
        engine_type = "TestEngine2"  # Use different name to avoid conflicts

        # Record enough data to create an estimate (need at least 3 attempts)
        for wait_time in [4.0, 5.0, 6.0]:
            self.tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=True,
                retry_count=1,
            )

        original_base = self.tracker.current_estimates[engine_type]["base"]

        # Create a new tracker instance (uses same database)
        new_tracker = AdaptiveRateLimitTracker()

        # Should load the previous estimate
        self.assertIn(engine_type, new_tracker.current_estimates)
        loaded_base = new_tracker.current_estimates[engine_type]["base"]

        # Should be close to the original (allowing for decay)
        self.assertAlmostEqual(loaded_base, original_base, delta=1.0)

        # Clean up
        new_tracker.reset_engine(engine_type)

    def test_get_stats(self):
        """Test getting statistics."""
        # Use unique engine name for this test to avoid conflicts
        import uuid

        engine_type = f"TestEngine_GetStats_{uuid.uuid4().hex[:8]}"

        # Create a fresh tracker for this test with programmatic_mode and rate limiting enabled
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"}
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # Make sure we start clean
        tracker.reset_engine(engine_type)

        # Record enough data to create an estimate (need at least 3 attempts)
        for wait_time in [3.0, 3.5, 4.0]:
            tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=True,
                retry_count=1,
            )

        # For get_stats, we'll check the in-memory estimates
        self.assertIn(engine_type, tracker.current_estimates)

        # Clean up
        tracker.reset_engine(engine_type)

    @pytest.mark.timeout(30)
    def test_reset_engine(self):
        """Test resetting an engine's data.

        Uses timeout marker instead of skipping to fail fast if DB operations hang.
        """
        # Use unique engine name for this test
        import uuid

        engine_type = f"TestEngine_Reset_{uuid.uuid4().hex[:8]}"

        # Use the test fixture's tracker instead of creating a new one
        # This ensures we're using the same database session
        tracker = self.tracker

        # Make sure we start clean
        tracker.reset_engine(engine_type)

        # Record enough data to create an estimate (need at least 3 attempts)
        for wait_time in [3.0, 3.5, 4.0]:
            tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=True,
                retry_count=1,
            )

        # Verify data exists
        self.assertIn(engine_type, tracker.current_estimates)

        # Get the wait time before reset - take multiple samples to account for randomness
        wait_times_before = [
            tracker.get_wait_time(engine_type) for _ in range(10)
        ]
        avg_wait_time_before = sum(wait_times_before) / len(wait_times_before)
        self.assertGreaterEqual(
            avg_wait_time_before, 2.5
        )  # Should be around recorded values

        # Reset the engine
        try:
            tracker.reset_engine(engine_type)
        except Exception as e:
            print(f"DEBUG: Exception during reset: {e}")
            # Even if database reset fails, memory should be cleared

        # Check immediately after reset, before calling get_wait_time
        # The engine should not be in estimates after reset
        self.assertNotIn(engine_type, tracker.current_estimates)

        # After reset, wait time should be much lower
        # Take multiple samples to account for randomness
        wait_times_after = [
            tracker.get_wait_time(engine_type) for _ in range(10)
        ]
        avg_wait_time_after = sum(wait_times_after) / len(wait_times_after)

        # After reset, should get default wait time (0.5s for unknown engines)
        # With some tolerance for CI environment variations
        self.assertLess(
            avg_wait_time_after, 1.0
        )  # Should be close to default 0.5s
        self.assertLess(
            avg_wait_time_after, avg_wait_time_before * 0.5
        )  # Should be significantly lower

    @pytest.mark.timeout(30)
    def test_reset_engine_simple(self):
        """Simple test for reset functionality.

        Uses timeout marker instead of skipping to fail fast if DB operations hang.
        """
        engine_type = "TestEngine_Simple_Reset"

        # Just test in-memory operations without database
        self.tracker.current_estimates[engine_type] = {
            "base": 5.0,
            "min": 2.0,
            "max": 10.0,
            "confidence": 0.8,
        }
        from collections import deque

        self.tracker.recent_attempts[engine_type] = deque(maxlen=100)

        # Verify data exists
        self.assertIn(engine_type, self.tracker.current_estimates)

        # Clear in-memory data directly (bypass database)
        if engine_type in self.tracker.current_estimates:
            del self.tracker.current_estimates[engine_type]
        if engine_type in self.tracker.recent_attempts:
            del self.tracker.recent_attempts[engine_type]

        # Verify data is cleared
        self.assertNotIn(engine_type, self.tracker.current_estimates)
        self.assertNotIn(engine_type, self.tracker.recent_attempts)

    def test_exploration_vs_exploitation(self):
        """Test that exploration sometimes returns different wait times."""
        engine_type = "TestEngine"

        # Set up a known estimate
        self.tracker.current_estimates[engine_type] = {
            "base": 10.0,
            "min": 5.0,
            "max": 20.0,
            "confidence": 0.8,
        }

        # Get multiple wait times
        wait_times = [
            self.tracker.get_wait_time(engine_type) for _ in range(50)
        ]

        # Should have some variation due to exploration and jitter
        unique_times = set(wait_times)
        self.assertGreater(len(unique_times), 5)  # Should have some variation

        # All should be within bounds
        for wait_time in wait_times:
            self.assertGreaterEqual(wait_time, 5.0)
            self.assertLessEqual(wait_time, 20.0)


class TestRateLimitIntegration(unittest.TestCase):
    """Test rate limiting integration with search engines."""

    def test_rate_limit_error_exception(self):
        """Test that RateLimitError can be raised and caught."""
        with self.assertRaises(RateLimitError):
            raise RateLimitError("Test rate limit")

    @patch(
        "local_deep_research.web_search_engines.rate_limiting.tracker.AdaptiveRateLimitTracker"
    )
    def test_base_search_engine_integration(self, mock_tracker_class):
        """Test integration with BaseSearchEngine."""
        # This would require more complex mocking of the search engine
        # For now, just verify the import works
        from local_deep_research.web_search_engines.search_engine_base import (
            BaseSearchEngine,
        )

        # Create a mock engine to verify rate_tracker is set during init
        # We need to provide required abstract methods
        class MockSearchEngine(BaseSearchEngine):
            def _get_previews(self, query):
                return []

            def _get_full_content(self, relevant_items):
                return []

        # Create instance and verify rate_tracker is set
        mock_engine = MockSearchEngine()
        self.assertTrue(hasattr(mock_engine, "rate_tracker"))
        self.assertIsNotNone(mock_engine.rate_tracker)


class TestAdaptiveRateLimitTrackerProfiles(unittest.TestCase):
    """Test rate limiting profile configurations."""

    def test_conservative_profile(self):
        """Test conservative profile applies lower rates."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
            "rate_limiting.profile": {
                "value": "conservative",
                "ui_element": "select",
            },
            "rate_limiting.exploration_rate": {
                "value": 0.1,
                "ui_element": "slider",
            },
            "rate_limiting.learning_rate": {
                "value": 0.3,
                "ui_element": "slider",
            },
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # Conservative profile should reduce exploration rate
        self.assertLessEqual(tracker.exploration_rate, 0.05)
        # Conservative profile should reduce learning rate
        self.assertLessEqual(tracker.learning_rate, 0.21)

    def test_aggressive_profile(self):
        """Test aggressive profile applies higher rates."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
            "rate_limiting.profile": {
                "value": "aggressive",
                "ui_element": "select",
            },
            "rate_limiting.exploration_rate": {
                "value": 0.1,
                "ui_element": "slider",
            },
            "rate_limiting.learning_rate": {
                "value": 0.3,
                "ui_element": "slider",
            },
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # Aggressive profile should increase exploration rate
        self.assertGreaterEqual(tracker.exploration_rate, 0.1)
        # Aggressive profile should increase learning rate
        self.assertGreaterEqual(tracker.learning_rate, 0.3)

    def test_balanced_profile(self):
        """Test balanced profile keeps default rates."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
            "rate_limiting.profile": {
                "value": "balanced",
                "ui_element": "select",
            },
            "rate_limiting.exploration_rate": {
                "value": 0.1,
                "ui_element": "slider",
            },
            "rate_limiting.learning_rate": {
                "value": 0.3,
                "ui_element": "slider",
            },
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # Balanced profile should use configured rates
        self.assertEqual(tracker.exploration_rate, 0.1)
        self.assertEqual(tracker.learning_rate, 0.3)


class TestApplyRateLimit(unittest.TestCase):
    """Test the apply_rate_limit method."""

    def test_apply_rate_limit_disabled(self):
        """Test apply_rate_limit returns 0 when disabled."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": False, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        wait_time = tracker.apply_rate_limit("TestEngine")
        self.assertEqual(wait_time, 0.0)

    def test_apply_rate_limit_enabled(self):
        """Test apply_rate_limit returns and applies wait time when enabled."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # For unknown engines, returns default optimistic wait time
        wait_time = tracker.apply_rate_limit("TestEngine")
        self.assertGreaterEqual(wait_time, 0.0)


class TestUpdateEstimate(unittest.TestCase):
    """Test the _update_estimate method."""

    def test_update_estimate_needs_minimum_attempts(self):
        """Test that update_estimate requires at least 3 attempts."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        engine_type = "TestEngine_MinAttempts"

        # Record only 2 attempts
        tracker.record_outcome(
            engine_type=engine_type,
            wait_time=1.0,
            success=True,
            retry_count=1,
        )
        tracker.record_outcome(
            engine_type=engine_type,
            wait_time=1.5,
            success=True,
            retry_count=1,
        )

        # Should not have an estimate yet
        self.assertNotIn(engine_type, tracker.current_estimates)

        # Third attempt should create estimate
        tracker.record_outcome(
            engine_type=engine_type,
            wait_time=2.0,
            success=True,
            retry_count=1,
        )

        self.assertIn(engine_type, tracker.current_estimates)

    def test_update_estimate_all_failures(self):
        """Test that all failures increase wait time."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        engine_type = "TestEngine_AllFail"

        # Record 3 failed attempts
        for wait_time in [1.0, 1.5, 2.0]:
            tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=False,
                retry_count=1,
            )

        # Should have an estimate with increased base time
        self.assertIn(engine_type, tracker.current_estimates)
        # Base should be higher than max failed wait * 1.5, capped at 10
        self.assertGreater(tracker.current_estimates[engine_type]["base"], 2.0)


class TestGetStats(unittest.TestCase):
    """Test the get_stats method."""

    def test_get_stats_in_ci_mode(self):
        """Test get_stats returns in-memory stats in CI mode."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        engine_type = "TestEngine_Stats"

        # Record some attempts
        for wait_time in [1.0, 1.5, 2.0]:
            tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=True,
                retry_count=1,
            )

        # Mock is_ci_environment to return True, ensuring we take the in-memory path
        with patch(
            "local_deep_research.web_search_engines.rate_limiting.tracker.is_ci_environment",
            return_value=True,
        ):
            # Get stats
            stats = tracker.get_stats(engine_type)

            # Should return stats from in-memory estimates
            self.assertIsInstance(stats, list)
            self.assertEqual(len(stats), 1)

    def test_get_stats_all_engines_in_ci_mode(self):
        """Test get_stats with no engine type returns all in CI mode."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # Record attempts for multiple engines
        for engine in ["Engine1", "Engine2"]:
            for wait_time in [1.0, 1.5, 2.0]:
                tracker.record_outcome(
                    engine_type=engine,
                    wait_time=wait_time,
                    success=True,
                    retry_count=1,
                )

        # Mock is_ci_environment to return True
        with patch(
            "local_deep_research.web_search_engines.rate_limiting.tracker.is_ci_environment",
            return_value=True,
        ):
            # Get all stats
            stats = tracker.get_stats()

            # Should return stats for both engines
            self.assertGreaterEqual(len(stats), 2)


class TestRateLimitingDisabled(unittest.TestCase):
    """Test behavior when rate limiting is disabled."""

    def test_disabled_returns_minimal_wait(self):
        """Test disabled rate limiting returns minimal wait time."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": False, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        wait_time = tracker.get_wait_time("AnyEngine")
        self.assertEqual(wait_time, 0.1)

    def test_disabled_does_not_record_outcomes(self):
        """Test disabled rate limiting doesn't record outcomes."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": False, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        engine_type = "TestEngine_Disabled"

        # Try to record outcomes
        tracker.record_outcome(
            engine_type=engine_type,
            wait_time=1.0,
            success=True,
            retry_count=1,
        )

        # Should not create any tracking data
        self.assertNotIn(engine_type, tracker.current_estimates)
        self.assertNotIn(engine_type, tracker.recent_attempts)


class TestProgrammaticMode(unittest.TestCase):
    """Test programmatic mode behavior."""

    def test_programmatic_mode_skips_database(self):
        """Test programmatic mode skips database operations."""
        settings_snapshot = {
            "rate_limiting.enabled": {"value": True, "ui_element": "checkbox"},
        }
        tracker = AdaptiveRateLimitTracker(
            settings_snapshot=settings_snapshot, programmatic_mode=True
        )

        # Should work without database errors
        engine_type = "TestEngine_Programmatic"

        for wait_time in [1.0, 1.5, 2.0]:
            tracker.record_outcome(
                engine_type=engine_type,
                wait_time=wait_time,
                success=True,
                retry_count=1,
            )

        # Should have in-memory estimate
        self.assertIn(engine_type, tracker.current_estimates)

    def test_programmatic_mode_default_disabled(self):
        """Test programmatic mode defaults to disabled rate limiting."""
        tracker = AdaptiveRateLimitTracker(programmatic_mode=True)

        # Without explicit settings, programmatic mode disables rate limiting
        self.assertFalse(tracker.enabled)


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