Locust load test script with custom load shaping.

SampleApp.py

import os
from locust import HttpUser, LoadTestShape, TaskSet, between, task

class UserTasks(TaskSet):
    """
    Represents the user tasks for load testing (not sequential).

    This class defines the tasks that the user will perform during the load testing.

    It is functionally similar to the JMeter test plan (SampleApp.jmx), where the user performs a series of tasks to simulate real-world
    """

    @task
    def lasttimestamp(self):
        """
        Performs a GET request to retrieve the timestamp of the last visit to the site.

        This method sends a GET request to the "/lasttimestamp" endpoint.
        """
        self.client.get("/lasttimestamp")

    @task
    def add(self):
        """
        Performs a POST request to increment the visitor counter by 1.

        This method sends a POST request to the "/add" endpoint with data parameter set to "1".
        """
        self.client.post("/add", data="1")

    @task
    def get(self):
        """
        Performs a GET request to retrieve the current the visitor count.

        This method sends a GET request to the "/get" endpoint.
        """
        self.client.get("/get")


class WebsiteUser(HttpUser):
    wait_time = between(0.5, 2) # Represents the wait time between consecutive tasks performed by the user.
    tasks = [UserTasks]

class StagesShape(LoadTestShape):
    """
    A simply load test shape class that has different user and spawn_rate at
    different stages.

    Keyword arguments:

        stages -- A list of dicts, each representing a stage with the following keys:
            duration -- When this many seconds pass the test is advanced to the next stage
            users -- Total user count
            spawn_rate -- Number of users to start/stop per second
            stop -- A boolean that can stop that test at a specific stage

        stop_at_end -- Can be set to stop once all stages have run.
    """

    users = os.getenv("LOCUST_USERS", 10) # Total number of users to simulate
    runtime = os.getenv("LOCUST_RUN_TIME", 180) # Total runtime of the test
    rampup = runtime * 0.1 # Ramp-up time for the test (10% of total runtime)
    stable = runtime * 0.85 # Stable time for the test (85% of total runtime)
    rampdown = runtime * 0.05 # Ramp-down time for the test (5% of total runtime)

    # Set minimum test duration for each stage (1 min in total)
    if rampup < 10:
        rampup = 10
    if stable < 40:
        stable = 40
    if rampdown < 10:
        rampdown = 10

    span_rate = users/rampup

    stages = [
        {"duration": rampup, "users": users, "spawn_rate": span_rate},  # Ramp-up linearly to max users during ramp up
        {"duration": stable, "users": users, "spawn_rate": 1},          # Stay at max users for main part of test
        {"duration": rampdown, "users": 1, "spawn_rate": 1},            # Linearly decrease to 1 user for remainder of test
    ]

    stop_at_end = True # Setting this to True will stop the test once all stages have run

    def tick(self):
        run_time = self.get_run_time()

        for stage in self.stages:
            if run_time < stage["duration"]:
                tick_data = (stage["users"], stage["spawn_rate"])
                return tick_data

        return None