def run_load_tests()

in right_size_your_sagemaker_endpoints/load_test_helper.py [0:0]


def run_load_tests(api_url, endpoints_list):
    """
    Runs a locust load test for a given list of endpoints.
    
    Given a set of SageMaker endpoints and an API Gateway URL, this
    function runs the run_locust.sh file for each endpoint. For easy 
    reference, the result files are organized into a folder with the 
    endpoint name and time prefix.
    
    Inputs:
    api_url: API Gateway URL for load testing.
    endpoints_list: List of SageMaker endpoints to run the load tests on.
    
    Output:
    Folder name where the load test results are saved.
    """ 
    
    time_stamp = strftime("%Y-%m-%d-%H-%M-%S", gmtime())
    # Iterate over the different endpoints
    for endpoint in endpoints_list:
        
        print(f"\nLoad testing {endpoint}...")
        subprocess.run(["bash", "run_locust.sh", f"{endpoint}", f"{api_url}"], stdout=subprocess.PIPE)
        
        print(f"Organizing {endpoint} result files...")
        
        prefix = f"test_{endpoint}"
        os.makedirs(f"results-{time_stamp}/{prefix}")
        
        files = [f for f in os.listdir() if prefix in f]
        for file in files:
            os.rename(file, f"results-{time_stamp}/{prefix}/{file}")
            
    print("Load testing complete!")
            
    return f"results-{time_stamp}"