radionica3d/backend/scratch/load_test.py

import asyncio
import httpx
import websockets
import json
import time
import random

BASE_URL = "http://localhost:8000"
WS_URL = "ws://localhost:8000"

async def simulate_user(user_id):
    """Simulates a single user behavior"""
    email = f"testuser_{user_id}@example.com"
    password = "password123"
    
    async with httpx.AsyncClient() as client:
        # 0. Try to register first (in case users don't exist)
        try:
            await client.post(f"{BASE_URL}/auth/register", json={
                "email": email,
                "password": password,
                "first_name": "Test",
                "last_name": "User",
                "phone": "123456"
            })
        except: pass

        # 1. Try to login
        print(f"[User {user_id}] Attempting login...")
        for i in range(5):
            try:
                resp = await client.post(f"{BASE_URL}/auth/login", json={
                    "email": email,
                    "password": password
                })
                if resp.status_code == 200:
                    token = resp.json().get("access_token")
                    print(f"[User {user_id}] Login successful on attempt {i+1}")
                    break
                elif resp.status_code == 403:
                    print(f"[User {user_id}] Login failed: Captcha required or rate limited")
                    # If captcha is required, this user stops login attempts in this simple script
                    return 
            except Exception as e:
                print(f"[User {user_id}] Login error: {e}")
                return
        else:
            print(f"[User {user_id}] Failed to login after 5 attempts")
            return

        # 2. Connect to global websocket
        ws_uri = f"{WS_URL}/auth/ws/global?token={token}"
        try:
            async with websockets.connect(ws_uri) as ws:
                print(f"[User {user_id}] WebSocket connected")
                
                # 3. Simulate order creation (to test 1-min rate limit)
                # We try twice rapidly
                for i in range(2):
                    order_data = {
                        "first_name": "Test",
                        "last_name": "User",
                        "phone": "123456",
                        "email": email,
                        "shipping_address": "Test Load Street",
                        "material_id": 1,
                        "quantity": 1
                    }
                    # Need to use form-data with file
                    files = {'file': ('model.stl', b'binary content stl', 'application/octet-stream')}
                    resp = await client.post(f"{BASE_URL}/orders", data=order_data, files=files)
                    
                    if resp.status_code == 200:
                        print(f"[User {user_id}] Order {i+1} created successfully")
                    elif resp.status_code == 429:
                        print(f"[User {user_id}] Order {i+1} REJECTED (Rate Limit Working!)")
                    else:
                        print(f"[User {user_id}] Order {i+1} failed: {resp.status_code}")
                
                # 4. Stay connected for a bit to receive potential updates
                # and send pings
                for _ in range(5):
                    await ws.send("ping")
                    try:
                        # Wait for a message with timeout
                        msg = await asyncio.wait_for(ws.recv(), timeout=2.0)
                        print(f"[User {user_id}] WS Received: {msg[:50]}...")
                    except asyncio.TimeoutError:
                        pass
                    await asyncio.sleep(1)
                    
        except Exception as e:
            print(f"[User {user_id}] WS/Order error: {e}")

async def main():
    print("--- Starting Load Test ---")
    start_time = time.time()
    
    # Run 10 concurrent users
    tasks = [simulate_user(i) for i in range(10)]
    await asyncio.gather(*tasks)
    
    end_time = time.time()
    print(f"--- Load Test Finished in {end_time - start_time:.2f}s ---")

if __name__ == "__main__":
    asyncio.run(main())