locustfile.py

from locust import HttpUser, task, between
from dotenv import load_dotenv
import os
import hashlib
import random

load_dotenv()
HOST = os.getenv('host')
DEMO_SECRET = os.getenv('x-demo-secret')


class APIUser(HttpUser):
    wait_time = between(1, 3)  # Wait 1-3 seconds between tasks
    host = HOST
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # You'll need to set these values in your environment
        self.headers = {
            "Content-Type": "application/json",
            "x-demo-secret": DEMO_SECRET
        }

    @task(2)  # Higher weight for image generation tasks
    def test_draw_endpoints(self):
        prompts = [
            "A friendly cartoon bear with expressive eyes, front view portrait",
            "A wise old wizard with a long white beard, looking directly at viewer",
            "A cheerful robot with glowing blue eyes, head and shoulders portrait",
            "A mystical elf warrior with pointed ears and silver hair, 3/4 view",
            "A cute animated fox with orange fur and big eyes, front facing",
            "A steampunk inventor with brass goggles and wild hair, portrait",
            "A gentle dragon with iridescent scales, close-up face shot",
            "A ferocious snake with a long, forked tongue, close-up face shot",
        ]
        for prompt in prompts:
            self.client.post("/test-draw", 
                           json={"prompt": prompt},
                           headers=self.headers)

    @task(7)
    def test_chat_endpoints(self):
        questions = [
            "Hi",
            "How are you?",
            "What is your name?",
            "What is your favorite color?",
            "What is your favorite food?",
            "What is your favorite animal?",
            "What is your favorite movie?",
            "What is your favorite book?",
            "What is your favorite song?",
        ]
        
        characters = [
            "A friendly cartoon bear with expressive eyes, front view portrait",
            "A wise old wizard with a long white beard, looking directly at viewer",
            "A cheerful robot with glowing blue eyes, head and shoulders portrait",
            "A mystical elf warrior with pointed ears and silver hair, 3/4 view",
            "A cute animated fox with orange fur and big eyes, front facing",
            "A steampunk inventor with brass goggles and wild hair, portrait",
            "A gentle dragon with iridescent scales, close-up face shot",
            "A ferocious snake with a long, forked tongue, close-up face shot",
        ]
        character = random.choice(characters)
        system_prompt = f"Imagine you are {character}"

        # Generate consistent session ID based on character
        session_id = f'locust-{hashlib.md5(character.encode()).hexdigest()}'
        
        for question in questions:
            self.client.post("/test-chat",
                           json={
                               "system_prompt": system_prompt,
                               "question": question,
                               "session_id": session_id
                           },
                           headers=self.headers)

    @task(1)  # Lower weight since this is a two-step process
    def test_animation_flow(self):
        # Step 1: Start animation
        response = self.client.post("/test-animate",
                                  json={"text": "The animated bear looks at the camera and says 'What's up?'"},
                                  headers=self.headers)
        
        # Step 2: Check status (if the first request was successful)
        if response.status_code == 200:
            invocation_id = response.json().get("invocationId")
            if invocation_id:
                self.client.get(f"/check-clip-status/{invocation_id}",
                              headers=self.headers)