"""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