from locust import task, HttpUser
import random
class APIKeyTestUser(HttpUser):
def on_start(self):
# API keys for testing (replace with your actual keys)
self.valid_key = "sk-1234567890abcdef"
self.invalid_key = "invalid-key-12345"
# API endpoints that require authentication
self.endpoints = ["/api/users", "/api/data", "/api/reports"]
@task(4)
def test_header_api_key(self):
"""Test API key in header (X-API-Key: your-key)"""
endpoint = random.choice(self.endpoints)
headers = {
"X-API-Key": self.valid_key,
"Content-Type": "application/json"
}
with self.client.get(
endpoint,
headers=headers,
name="Header API Key"
) as response:
if response.status_code == 200:
print(f"Header API key success: {endpoint}")
elif response.status_code == 401:
response.failure("Valid API key rejected")
else:
response.failure(f"Unexpected response: {response.status_code}")
@task(3)
def test_bearer_token_api_key(self):
"""Test API key as Bearer token (Authorization: Bearer your-key)"""
endpoint = random.choice(self.endpoints)
headers = {
"Authorization": f"Bearer {self.valid_key}",
"Content-Type": "application/json"
}
with self.client.get(
endpoint,
headers=headers,
name="Bearer Token API Key"
) as response:
if response.status_code == 200:
print(f"Bearer token success: {endpoint}")
elif response.status_code == 401:
response.failure("Valid bearer token rejected")
else:
response.failure(f"Unexpected response: {response.status_code}")
@task(3)
def test_query_parameter_api_key(self):
"""Test API key in query parameter (?api_key=your-key)"""
endpoint = random.choice(self.endpoints)
params = {"api_key": self.valid_key}
with self.client.get(
endpoint,
params=params,
name="Query Parameter API Key"
) as response:
if response.status_code == 200:
print(f"Query parameter success: {endpoint}")
elif response.status_code == 401:
response.failure("Valid API key in query rejected")
else:
response.failure(f"Unexpected response: {response.status_code}")
@task(2)
def test_invalid_api_key(self):
"""Test with invalid API key"""
endpoint = random.choice(self.endpoints)
headers = {
"X-API-Key": self.invalid_key,
"Content-Type": "application/json"
}
with self.client.get(
endpoint,
headers=headers,
name="Invalid API Key"
) as response:
if response.status_code == 401:
print(f"Invalid API key correctly rejected: {endpoint}")
elif response.status_code == 200:
response.failure("Invalid API key was accepted")
else:
print(f"Invalid API key returned {response.status_code}")
@task(1)
def test_missing_api_key(self):
"""Test request without API key"""
endpoint = random.choice(self.endpoints)
with self.client.get(
endpoint,
name="Missing API Key"
) as response:
if response.status_code == 401:
print(f"Missing API key correctly rejected: {endpoint}")
elif response.status_code == 200:
response.failure("Request without API key was accepted")
else:
print(f"Missing API key returned {response.status_code}")