locust_boto3_custom_client

locust_boto3_custom_client.py

import boto3
import json
from locust import events, Locust, task, TaskSet
import time


class SimpleSentTask(TaskSet):
    @task
    def send_request(self):
        payload = {"sample_key": "sample_value"}
        self.client.send("send_endpoint", "POST", payload, "send_endpoint_task_name")

class SimpleSendRequest(Locust):
    def __init__ (self, *args, **kwargs):
        super(Locust, self).__init__(*args, **kwargs)
        self.client = BotoClient()

    task_set = SimpleSentTask
    min_wait = 5000
    max_wait = 15000


class BotoClient:
    def __init__(self):
        self.client = boto3.client("apigateway" , region_name="eu-west-2")
        self.api_resource_ids = get_resource_ids('your-api-name')

    def get_api_resource_ids(self, api_gateway_id: str) -> dict:
        resource_ids = {}
        for resource in self.client.get_resources(restApiId=api_gateway_id)["items"]:
            try:
                resource_ids[resource["pathPart"]] = resource["id"]
            except KeyError:
                pass # Ignore resources with no pathPart
            return resource_ids

    def get_resource_ids(self, api_gateway_name: str) -> dict:
        api_gateways = self.client.get_rest_apis()["items"]
        for api_gateway in api_gateways:
            if api_gateway["name"] == api_gateway_name:
                return get_api_resource_ids(api_gateway["id"])
        raise ValueError(f"{api_name} not found.")

    @staticmethod
    def total_time(start_time) -> float:
        return int((time.time() - start_time) * 1000)

    def send(self, name: str, http_method: str, payload: dict, endpoint_name: str):
        start_time = time.time()
        try:
            self.client.test_invoke_method(
                restApiId=self.api_resource_ids["api"],
                resourceId=self.api_resource_ids[endpoint_name],
                httpMethod=http_method,
                headers={"Content-Type": "application/json"},
                body=json.dumps(payload)
            )
        except Exception as e:
            events.request_failure.fire(
                request_type="execute",
                name=name,
                response_time=total_time(start_time),
                exception=e
            )
        events.request_success.fire(
            request_type="execute",
            name=name,
            response_time=total_time(start_time),
            response_length=0
        )