Simple Rate Limiting

from locust import task, HttpUser
import time
import random

class RateLimitUser(HttpUser):
    def on_start(self):
        self.backoff_until = 0  # Track when to resume requests
        self.rate_limit_count = 0

    @task(5)
    def test_api_with_rate_limits(self):
        """Test API endpoint that has rate limits"""
        
        # Check if we're in backoff period
        if time.time() < self.backoff_until:
            print("Skipping request - in backoff period")
            return
        
        with self.client.get("/api/data", name="Rate Limited API") as response:
            if response.status_code == 200:
                print("Request successful")
                
                # Check rate limit headers
                remaining = response.headers.get("X-RateLimit-Remaining")
                if remaining:
                    print(f"Requests remaining: {remaining}")
                    
                    # Warn if getting close to limit
                    if int(remaining) < 5:
                        print("WARNING: Approaching rate limit!")
                        
            elif response.status_code == 429:
                # Rate limited - implement backoff
                print("Rate limited! Backing off...")
                self.rate_limit_count += 1
                
                # Get retry-after header or use default
                retry_after = response.headers.get("Retry-After", "60")
                backoff_time = int(retry_after)
                
                # Set backoff period
                self.backoff_until = time.time() + backoff_time
                print(f"Backing off for {backoff_time} seconds")
                
                response.failure(f"Rate limited - backing off for {backoff_time}s")
            else:
                response.failure(f"Request failed: {response.status_code}")

    @task(2)
    def test_search_api(self):
        """Test search API with rate limits"""
        
        if time.time() < self.backoff_until:
            return
        
        search_terms = ["test", "product", "user", "data"]
        query = random.choice(search_terms)
        
        with self.client.get(f"/api/search?q={query}", name="Search API") as response:
            if response.status_code == 200:
                print(f"Search for '{query}' successful")
            elif response.status_code == 429:
                print(f"Search rate limited")
                
                # Simple backoff
                self.backoff_until = time.time() + 30  # 30 second backoff
                response.failure("Search rate limited")
            else:
                response.failure(f"Search failed: {response.status_code}")

    @task(1)
    def check_rate_limit_status(self):
        """Check current rate limit status"""
        
        if time.time() < self.backoff_until:
            remaining_backoff = self.backoff_until - time.time()
            print(f"In backoff for {remaining_backoff:.0f} more seconds")
        else:
            print("No active backoff - ready for requests")
            
        if self.rate_limit_count > 0:
            print(f"Total rate limits hit: {self.rate_limit_count}")