learningcircuit--local-deep-research/tests/database/test_rate_limiting_models.py

"""Tests for rate limiting database models."""

import time

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from local_deep_research.database.models import (
    Base,
    RateLimitAttempt,
    RateLimitEstimate,
)


class TestRateLimitingModels:
    """Test suite for rate limiting models."""

    @pytest.fixture
    def engine(self):
        """Create an in-memory SQLite database for testing."""
        engine = create_engine("sqlite:///:memory:")
        Base.metadata.create_all(engine)
        yield engine
        engine.dispose()

    @pytest.fixture
    def session(self, engine):
        """Create a database session for testing."""
        Session = sessionmaker(bind=engine)
        session = Session()
        yield session
        session.close()

    def test_rate_limit_attempt_creation(self, session):
        """Test creating rate limit attempt records."""
        attempt = RateLimitAttempt(
            engine_type="google",
            timestamp=time.time(),
            wait_time=2.5,
            retry_count=1,
            success=True,
            error_type=None,
        )

        session.add(attempt)
        session.commit()

        # Verify attempt
        saved = session.query(RateLimitAttempt).first()
        assert saved is not None
        assert saved.engine_type == "google"
        assert saved.wait_time == 2.5
        assert saved.retry_count == 1
        assert saved.success is True
        assert saved.error_type is None

    def test_failed_rate_limit_attempts(self, session):
        """Test tracking failed rate limit attempts."""
        current_time = time.time()

        # Successful attempt
        success = RateLimitAttempt(
            engine_type="bing",
            timestamp=current_time,
            wait_time=1.0,
            retry_count=0,
            success=True,
        )

        # Failed attempt - too fast
        too_fast = RateLimitAttempt(
            engine_type="bing",
            timestamp=current_time + 1,
            wait_time=0.5,
            retry_count=1,
            success=False,
            error_type="rate_limit",
        )

        # Failed attempt - other error
        other_error = RateLimitAttempt(
            engine_type="bing",
            timestamp=current_time + 2,
            wait_time=2.0,
            retry_count=2,
            success=False,
            error_type="connection",
        )

        session.add_all([success, too_fast, other_error])
        session.commit()

        # Analyze attempts
        all_attempts = (
            session.query(RateLimitAttempt).filter_by(engine_type="bing").all()
        )
        assert len(all_attempts) == 3

        failed = session.query(RateLimitAttempt).filter_by(success=False).all()
        assert len(failed) == 2

        # Check error types
        rate_limit_errors = (
            session.query(RateLimitAttempt)
            .filter_by(error_type="rate_limit")
            .count()
        )
        assert rate_limit_errors == 1

    def test_rate_limit_estimate(self, session):
        """Test rate limit estimate storage and updates."""
        estimate = RateLimitEstimate(
            engine_type="duckduckgo",
            base_wait_seconds=1.5,
            min_wait_seconds=0.5,
            max_wait_seconds=10.0,
            last_updated=time.time(),
            total_attempts=100,
            success_rate=0.85,
        )

        session.add(estimate)
        session.commit()

        # Verify estimate
        saved = (
            session.query(RateLimitEstimate)
            .filter_by(engine_type="duckduckgo")
            .first()
        )
        assert saved is not None
        assert saved.base_wait_seconds == 1.5
        assert saved.success_rate == 0.85
        assert saved.total_attempts == 100

    def test_multiple_engine_tracking(self, session):
        """Test tracking rate limits for multiple search engines."""
        engines = [
            ("google", 1.0, 0.5, 5.0, 0.9),
            ("bing", 0.8, 0.3, 3.0, 0.92),
            ("duckduckgo", 0.5, 0.1, 2.0, 0.95),
            ("searx", 2.0, 1.0, 10.0, 0.8),
        ]

        current_time = time.time()

        for engine, base, min_val, max_val, success_rate in engines:
            estimate = RateLimitEstimate(
                engine_type=engine,
                base_wait_seconds=base,
                min_wait_seconds=min_val,
                max_wait_seconds=max_val,
                last_updated=current_time,
                total_attempts=50,
                success_rate=success_rate,
            )
            session.add(estimate)

        session.commit()

        # Verify all engines
        all_estimates = session.query(RateLimitEstimate).all()
        assert len(all_estimates) == 4

        # Find most reliable engine
        most_reliable = (
            session.query(RateLimitEstimate)
            .order_by(RateLimitEstimate.success_rate.desc())
            .first()
        )
        assert most_reliable.engine_type == "duckduckgo"

    def test_adaptive_rate_learning(self, session):
        """Test updating rate limit estimates based on attempts."""
        # Initial estimate
        estimate = RateLimitEstimate(
            engine_type="adaptive_test",
            base_wait_seconds=1.0,
            min_wait_seconds=0.5,
            max_wait_seconds=5.0,
            last_updated=time.time(),
            total_attempts=10,
            success_rate=0.8,
        )

        session.add(estimate)
        session.commit()

        # Simulate attempts
        current_time = time.time()
        attempts = [
            (0.5, False),  # Too fast
            (0.8, False),  # Still too fast
            (1.2, True),  # Success
            (1.1, True),  # Success
            (1.0, True),  # Success
            (0.9, False),  # Too fast again
            (1.1, True),  # Success
        ]

        for i, (wait_time, success) in enumerate(attempts):
            attempt = RateLimitAttempt(
                engine_type="adaptive_test",
                timestamp=current_time + i,
                wait_time=wait_time,
                retry_count=0 if success else 1,
                success=success,
                error_type=None if success else "rate_limit",
            )
            session.add(attempt)

        session.commit()

        # Update estimate based on attempts
        successful_waits = (
            session.query(RateLimitAttempt.wait_time)
            .filter_by(engine_type="adaptive_test", success=True)
            .all()
        )

        if successful_waits:
            avg_successful_wait = sum(w[0] for w in successful_waits) / len(
                successful_waits
            )
            estimate.base_wait_seconds = avg_successful_wait
            estimate.total_attempts += len(attempts)
            estimate.success_rate = len(successful_waits) / len(attempts)
            estimate.last_updated = time.time()
            session.commit()

        # Verify updated estimate
        updated = (
            session.query(RateLimitEstimate)
            .filter_by(engine_type="adaptive_test")
            .first()
        )
        assert (
            updated.base_wait_seconds > 1.0
        )  # Should increase based on attempts
        assert updated.total_attempts == 17  # 10 + 7
        assert updated.success_rate == 4 / 7  # 4 successes out of 7 attempts

    def test_time_based_patterns(self, session):
        """Test identifying time-based rate limit patterns."""
        # Simulate different success rates at different times
        base_time = time.time()

        # Morning hours - higher success rate
        for i in range(10):
            attempt = RateLimitAttempt(
                engine_type="time_pattern",
                timestamp=base_time + i * 60,  # Every minute
                wait_time=1.0,
                retry_count=0,
                success=i % 3 != 0,  # 66% success
                error_type=None if i % 3 != 0 else "rate_limit",
            )
            session.add(attempt)

        # Afternoon - lower success rate
        for i in range(10):
            attempt = RateLimitAttempt(
                engine_type="time_pattern",
                timestamp=base_time + 3600 + i * 60,  # 1 hour later
                wait_time=1.0,
                retry_count=0,
                success=i % 2 == 0,  # 50% success
                error_type=None if i % 2 == 0 else "rate_limit",
            )
            session.add(attempt)

        session.commit()

        # Analyze patterns
        morning_attempts = (
            session.query(RateLimitAttempt)
            .filter(
                RateLimitAttempt.engine_type == "time_pattern",
                RateLimitAttempt.timestamp < base_time + 3600,
            )
            .all()
        )

        afternoon_attempts = (
            session.query(RateLimitAttempt)
            .filter(
                RateLimitAttempt.engine_type == "time_pattern",
                RateLimitAttempt.timestamp >= base_time + 3600,
            )
            .all()
        )

        morning_success_rate = sum(
            1 for a in morning_attempts if a.success
        ) / len(morning_attempts)
        afternoon_success_rate = sum(
            1 for a in afternoon_attempts if a.success
        ) / len(afternoon_attempts)

        assert morning_success_rate > afternoon_success_rate

    def test_estimate_updates(self, session):
        """Test updating rate limit estimates."""
        # Create initial estimate
        estimate = RateLimitEstimate(
            engine_type="update_test",
            base_wait_seconds=1.0,
            min_wait_seconds=0.5,
            max_wait_seconds=5.0,
            last_updated=time.time() - 3600,  # 1 hour ago
            total_attempts=50,
            success_rate=0.8,
        )

        session.add(estimate)
        session.commit()

        # Update estimate
        estimate.base_wait_seconds = 1.5
        estimate.success_rate = 0.85
        estimate.total_attempts = 75
        estimate.last_updated = time.time()

        session.commit()

        # Verify updates
        updated = (
            session.query(RateLimitEstimate)
            .filter_by(engine_type="update_test")
            .first()
        )
        assert updated.base_wait_seconds == 1.5
        assert updated.success_rate == 0.85
        assert updated.total_attempts == 75
        assert updated.last_updated > time.time() - 60  # Updated recently

    def test_cleanup_old_attempts(self, session):
        """Test cleaning up old rate limit attempts."""
        current_time = time.time()

        # Create attempts at different ages
        for days_ago in range(10):
            attempt = RateLimitAttempt(
                engine_type="cleanup_test",
                timestamp=current_time - (days_ago * 86400),  # Days in seconds
                wait_time=1.0,
                retry_count=0,
                success=True,
            )
            session.add(attempt)

        session.commit()

        # Count old attempts (older than 7 days)
        seven_days_ago = current_time - (7 * 86400)
        old_attempts = (
            session.query(RateLimitAttempt)
            .filter(RateLimitAttempt.timestamp < seven_days_ago)
            .count()
        )

        # For days_ago in [7, 8, 9], which are all > 7 days ago
        # But since timestamp is current_time - (days_ago * 86400)
        # Only days 8 and 9 are actually older than 7 days
        assert old_attempts == 2  # Days 8, 9

        # Delete old attempts
        session.query(RateLimitAttempt).filter(
            RateLimitAttempt.timestamp < seven_days_ago
        ).delete()
        session.commit()

        # Verify cleanup
        remaining = session.query(RateLimitAttempt).count()
        assert remaining == 8  # 10 total - 2 deleted = 8

    def test_rate_limit_metadata(self, session):
        """Test storing metadata with attempts."""
        attempt = RateLimitAttempt(
            engine_type="metadata_test",
            timestamp=time.time(),
            wait_time=2.0,
            retry_count=1,
            success=True,
            error_type=None,
        )

        session.add(attempt)
        session.commit()

        # Verify
        saved = session.query(RateLimitAttempt).first()
        assert saved.engine_type == "metadata_test"
        assert saved.created_at is not None  # Auto-populated

    def test_concurrent_engine_limits(self, session):
        """Test tracking concurrent rate limits for multiple engines."""
        current_time = time.time()

        # Create attempts for multiple engines at the same time
        engines = ["google", "bing", "duckduckgo"]

        for engine in engines:
            for i in range(5):
                attempt = RateLimitAttempt(
                    engine_type=engine,
                    timestamp=current_time + i,
                    wait_time=1.0 + i * 0.1,
                    retry_count=0,
                    success=True,
                )
                session.add(attempt)

        session.commit()

        # Verify each engine has its own attempts
        for engine in engines:
            count = (
                session.query(RateLimitAttempt)
                .filter_by(engine_type=engine)
                .count()
            )
            assert count == 5

        # Get latest attempt per engine
        from sqlalchemy import func

        latest_per_engine = (
            session.query(
                RateLimitAttempt.engine_type,
                func.max(RateLimitAttempt.timestamp).label("latest"),
            )
            .group_by(RateLimitAttempt.engine_type)
            .all()
        )

        assert len(latest_per_engine) == 3