load_test.py

#!/usr/bin/env python3
"""
Load test for JWT authentication service
"""

import requests
import time
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import statistics

class LoadTester:
    def __init__(self, base_url="http://localhost:8080"):
        self.base_url = base_url
        self.results = []
        
    def single_request(self, endpoint, method="GET", data=None, token=None):
        """Make a single request and return timing"""
        url = f"{self.base_url}{endpoint}"
        headers = {}
        
        if token:
            headers['Authorization'] = f'Bearer {token}'
        
        start_time = time.time()
        
        try:
            if method.upper() == "GET":
                response = requests.get(url, headers=headers, timeout=10)
            elif method.upper() == "POST":
                response = requests.post(url, json=data, headers=headers, timeout=10)
            else:
                return {"error": "Invalid method"}
            
            end_time = time.time()
            
            return {
                "endpoint": endpoint,
                "method": method,
                "status": response.status_code,
                "time": end_time - start_time,
                "success": response.status_code < 400
            }
            
        except Exception as e:
            end_time = time.time()
            return {
                "endpoint": endpoint,
                "method": method,
                "status": 0,
                "time": end_time - start_time,
                "success": False,
                "error": str(e)
            }
    
    def register_users(self, count=10):
        """Register multiple users for testing"""
        print(f"Registering {count} test users...")
        users = []
        
        for i in range(count):
            username = f"loadtest_user_{i}"
            password = f"password_{i}"
            
            data = {"username": username, "password": password}
            result = self.single_request("/api/auth/register", "POST", data)
            
            if result.get("status") in [201, 409]:  # Created or already exists
                users.append((username, password))
            
            time.sleep(0.1)  # Small delay between registrations
        
        print(f"Registered {len(users)} users")
        return users
    
    def test_endpoint(self, endpoint, method="GET", data=None, 
                      token=None, concurrent=10, requests_per_thread=10):
        """Test an endpoint with concurrent requests"""
        print(f"\nTesting {endpoint} with {concurrent} concurrent threads...")
        
        def worker():
            results = []
            for _ in range(requests_per_thread):
                result = self.single_request(endpoint, method, data, token)
                results.append(result)
            return results
        
        start_time = time.time()
        all_results = []
        
        with ThreadPoolExecutor(max_workers=concurrent) as executor:
            futures = [executor.submit(worker) for _ in range(concurrent)]
            
            for future in as_completed(futures):
                all_results.extend(future.result())
        
        end_time = time.time()
        
        # Analyze results
        successful = [r for r in all_results if r.get("success")]
        failed = [r for r in all_results if not r.get("success")]
        
        if successful:
            times = [r["time"] for r in successful]
            
            print(f"  Total requests: {len(all_results)}")
            print(f"  Successful: {len(successful)}")
            print(f"  Failed: {len(failed)}")
            print(f"  Total time: {end_time - start_time:.2f}s")
            print(f"  Requests per second: {len(all_results) / (end_time - start_time):.2f}")
            print(f"  Average response time: {statistics.mean(times):.3f}s")
            print(f"  Median response time: {statistics.median(times):.3f}s")
            print(f"  Min response time: {min(times):.3f}s")
            print(f"  Max response time: {max(times):.3f}s")
            
            if len(times) > 1:
                print(f"  Std deviation: {statistics.stdev(times):.3f}s")
        
        if failed:
            print(f"\n  Failed requests:")
            for fail in failed[:5]:  # Show first 5 failures
                print(f"    - {fail.get('error', 'Unknown error')}")
            if len(failed) > 5:
                print(f"    ... and {len(failed) - 5} more")
        
        return all_results
    
    def run_load_test(self, concurrent_users=10, requests_per_user=20):
        """Run comprehensive load test"""
        print(f"\n{'='*60}")
        print(f"LOAD TEST STARTING")
        print(f"Concurrent users: {concurrent_users}")
        print(f"Requests per user: {requests_per_user}")
        print(f"{'='*60}")
        
        # Step 1: Register test users
        test_users = self.register_users(concurrent_users)
        
        # Step 2: Test login endpoint
        print(f"\n{'='*60}")
        print("TEST 1: Login Endpoint Load Test")
        print(f"{'='*60}")
        
        login_tokens = []
        for username, password in test_users:
            data = {"username": username, "password": password}
            result = self.single_request("/api/auth/login", "POST", data)
            
            if result.get("success"):
                # In real test, we'd store the token
                login_tokens.append("dummy_token")
            
            time.sleep(0.05)
        
        # Step 3: Test login with concurrent requests
        self.test_endpoint(
            "/api/auth/login",
            "POST",
            {"username": test_users[0][0], "password": test_users[0][1]},
            concurrent=min(concurrent_users, 5),
            requests_per_thread=5
        )
        
        # Step 4: Test protected endpoint (would need actual tokens)
        print(f"\n{'='*60}")
        print("TEST 2: Protected Endpoint (simulated)")
        print(f"{'='*60}")
        print("Note: Need actual JWT tokens for proper protected endpoint testing")
        
        # Step 5: Test registration endpoint
        print(f"\n{'='*60}")
        print("TEST 3: Registration Endpoint Load Test")
        print(f"{'='*60}")
        
        self.test_endpoint(
            "/api/auth/register",
            "POST",
            {"username": "loadtest_new", "password": "password123"},
            concurrent=min(concurrent_users, 3),
            requests_per_thread=3
        )
        
        print(f"\n{'='*60}")
        print("LOAD TEST COMPLETED")
        print(f"{'='*60}")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Load test JWT authentication service')
    parser.add_argument('--url', default='http://localhost:8080', help='Base URL')
    parser.add_argument('--concurrent', type=int, default=5, help='Concurrent users')
    parser.add_argument('--requests', type=int, default=10, help='Requests per user')
    parser.add_argument('--quick', action='store_true', help='Run quick load test')
    
    args = parser.parse_args()
    
    tester = LoadTester(args.url)
    
    if args.quick:
        # Quick test
        print("Running quick load test...")
        tester.test_endpoint("/api/auth/login", "POST", 
                           {"username": "testuser", "password": "password123"},
                           concurrent=3, requests_per_thread=5)
    else:
        # Full load test
        tester.run_load_test(args.concurrent, args.requests)