Simple API Key Testing

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