LocustIO - sgml/signature GitHub Wiki

import os
import sys
from locust import HttpUser, task, between
from pathlib import Path

# Credentials from environment (sourced via .bashrc)
USERNAME = os.getenv("LOGIN_USERNAME")
PASSWORD = os.getenv("LOGIN_PASSWORD")

# Target base URL from command line argument
if len(sys.argv) < 2:
    print("Usage: locust -f script.py <base_url>")
    sys.exit(1)
BASE_URL = sys.argv[-1]

# Local path to save downloaded file
DOWNLOAD_PATH = Path("downloads/bop.csv")
DOWNLOAD_PATH.parent.mkdir(parents=True, exist_ok=True)

class WebsiteUser(HttpUser):
    host = BASE_URL
    wait_time = between(1, 2)

    def on_start(self):
        self.login()

    def login(self):
        with self.client.post(
            "/login",
            data={"username": USERNAME, "password": PASSWORD},
            catch_response=True
        ) as response:
            if response.status_code != 200:
                response.failure(f"Login failed: {response.status_code}")
                raise Exception(f"Login failed with status: {response.status_code}")
            response.success()

    @task
    def get_foo(self):
        self.check_status("/foo")

    @task
    def get_bar(self):
        self.check_status("/foo/bar")

    @task
    def get_baz(self):
        self.check_status("/foo/bar/baz")

    @task
    def download_csv(self):
        with self.client.get("/foo/bar/baz/bop.csv", stream=True, catch_response=True) as response:
            if response.status_code != 200:
                response.failure(f"Download failed: {response.status_code}")
                raise Exception(f"CSV download failed: {response.status_code}")
            with open(DOWNLOAD_PATH, "wb") as f:
                f.write(response.content)
            response.success()

        # Trigger logout request after successful download
        with self.client.get("/logout", catch_response=True) as logout_response:
            if logout_response.status_code != 200:
                logout_response.failure(f"Logout failed: {logout_response.status_code}")
                raise Exception(f"Logout failed with status: {logout_response.status_code}")
            logout_response.success()

'''
lr_jwt.py
'''

import jwt
import datetime

def long_running_jwt()
    # Define a secret key
    SECRET_KEY = "your_secret_key"

    # Create a payload with a long expiration time
    payload = {
        "user": "locust_test_user",
        "iat": datetime.datetime.utcnow(),  # Issued at
        "exp": datetime.datetime.utcnow() + datetime.timedelta(hours=6),  # Expires in 6 hours
    }

    # Generate the token
    token = jwt.encode(payload, SECRET_KEY, algorithm="HS256")

    print(f"Generated JWT: {token}")
    return token

from locust import HttpUser, task
from lr_jwt import long_running_jwt

# Replace with your generated JWT
JWT_TOKEN = long_running_jwt()

class LoadTestUser(HttpUser):
    @task
    def test_endpoint(self):
        headers = {"Authorization": f"Bearer {JWT_TOKEN}"}
        self.client.get("/your/api/endpoint", headers=headers)