boto3_class.py

import boto3
import json
from locust import events
import time

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
        )