[Vegeta + pdsh wrapper for distributed load testing] #distributed #loadtest #performance #vegeta #golang #python

Vegeta + pdsh wrapper for distributed load testing.py

#!/usr/bin/env python

"""Usage example:
./distributed-load-test.py 'GET https://foo.stage.example.com/beep/boop' -r 70 -d 30s -H newsiege01 newsiege02
"""

import argparse
import sys

from typing import List
from subprocess import run, Popen, PIPE


def parse_args():
    p = argparse.ArgumentParser(description='Distributed testing with vegeta and pdash')
    p.add_argument('target', help='vegeta targer, e.g. "GET https://httpbin.org/get"')
    p.add_argument('-H', '--hosts', action='store', type=str, nargs='*',
                   default=['newsiege01', 'newsiege02', 'newsiege03'], help='host list')
    p.add_argument('-r', '--rate', type=int, required=True,
                   help='Request rate, going to be devided evenly across hosts')
    p.add_argument('-d', '--duration', type=str, required=True,
                   help='Load test duration')
    p.add_argument('-e', '--executable', type=str, default='/usr/local/bin/vegeta',
                   help='Vegeta executable')
    p.add_argument('-o', '--output', default='stdout', type=str,
                   help='Output file')
    p.add_argument('-R', '--reporter', default='text', type=str,
                   help='Reporter [text, json, plot, hist[buckets]]')
    return p.parse_args()


def main(target: str, hosts: List[str], rate: int, duration: str, vegeta: str,
         output: str, reporter: str):
    cmd = f"pdsh -b -w '{','.join(hosts)}' 'echo \"{target}\" | vegeta attack -rate={int(rate/len(hosts))} -duration={duration} > result.bin'"
    pdsh = run(cmd, shell=True, stdout=PIPE, stderr=PIPE)
    if pdsh.returncode > 0:
        sys.stderr.write('pdsh failed:')
        sys.stderr.write(pdsh.stderr)
        sys.exit(pdsh.returncode)
    scps = []
    for host in hosts:
        scps.append(Popen(['scp', f'{host}:~/result.bin', f'{host}.bin'], stdout=PIPE, stderr=PIPE))
    [scp.communicate() for scp in scps]
    inputs = ','.join((f'{host}.bin' for host in hosts))
    run(f'{vegeta} report -inputs="{inputs}" -output={output} -reporter={reporter}', shell=True)


if __name__ == '__main__':
    args = parse_args()
    main(args.target, args.hosts, args.rate, args.duration, args.executable, args.output, args.reporter)