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