Locust MonKit wrapper for Jenkins

locust-wrapper.py

import argparse
import subprocess
import sys
import re
from collections import defaultdict

from lxml import etree

report = ["Min", "Avg", "Median", "95%", "99%"]

parser = argparse.ArgumentParser(usage="%s [options] [LocustClass [LocustClass2 ... ]]" % sys.argv[0])

parser.add_argument(
        '-H', '--host',
        dest="host",
        default=None,
        help="Host to load test in the following format: http://10.21.32.33"
)

parser.add_argument(
        '-f', '--locustfile',
        dest='locustfile',
        default='locustfile',
        help="Python module file to import, e.g. '../other.py'. Default: locustfile"
)

# Number of clients
parser.add_argument(
        '-c', '--clients',
        type=int,
        dest='num_clients',
        default=1,
        help="Number of concurrent clients."
)
# Client hatch rate
parser.add_argument(
        '-r', '--hatch-rate',
        type=float,
        dest='hatch_rate',
        default=1,
        help="The rate per second in which clients are spawned."
)
# Number of requests
parser.add_argument(
        '-n', '--num-request',
        type=int,
        dest='num_requests',
        required=True,
        help="Number of requests to perform."
)

parser.add_argument(
        'classes',
        nargs='*',
        help='Locust classes to run'
)

args = parser.parse_args()

locust_args = ["locust", "--no-web",
               "-f", args.locustfile,
               "-c", str(args.num_clients),
               "-r", str(args.hatch_rate),
               "-n", str(args.num_requests)
               ]
if args.host:
    locust_args.append("-H")
    locust_args.append(args.host)

locust_args += args.classes

process = subprocess.Popen(locust_args, stderr=subprocess.STDOUT, stdout=subprocess.PIPE)

result = ''

for line in iter(process.stdout.readline, ''):
    sys.stdout.write(line)

    if 'Shutting down (exit code ' in line or result:
        result += line

process.communicate()

if result:
    data = defaultdict(dict)

    match = re.search("Shutting down \(exit code \d\), bye\.\n.*?\n-+\n(.*?\n?)-+\n", result, re.DOTALL)
    if match:
        for row in match.group(1).split("\n"):
            if not row:
                continue

            row = row.strip()
            row = re.split("\s+", row)

            name = "%s %s" % (row[0], row[1])
            data[name]['reqs'] = row[2]
            data[name]['fails'] = row[3].split('(')[0]
            data[name]['Avg'] = row[4]
            data[name]['Min'] = row[5]
            data[name]['Max'] = row[6]
            data[name]['Median'] = row[8]
            data[name]['req/s'] = row[9]

    match = re.search("Percentage of the requests completed within given times\n.*?\n-+\n(.*?\n?)-+\n", result, re.DOTALL)
    if match:
        for row in match.group(1).split("\n"):
            if not row:
                continue

            row = row.strip()
            row = re.split("\s+", row)

            name = "%s %s" % (row[0], row[1])
            data[name]['50%'] = row[3].split('(')[0]
            data[name]['66%'] = row[4]
            data[name]['75%'] = row[5]
            data[name]['80%'] = row[6]
            data[name]['90%'] = row[7]
            data[name]['95%'] = row[8]
            data[name]['98%'] = row[9]
            data[name]['99%'] = row[10]
            data[name]['100%'] = row[11]

    root = etree.Element('categories')

    for row, rowdata in data.iteritems():
        category = etree.Element('category')
        category.set('name', row)
        category.set('scale', "ms")
        root.append(category)

        observations = etree.Element('observations')
        category.append(observations)

        for key in report:
            if key in rowdata:
                observation = etree.Element('observation')
                observation.set("name", key)
                observation.text = rowdata[key]
                observations.append(observation)

    if data:
        with open('monkit.xml', 'w+') as file:
            etree.ElementTree(root).write(file, pretty_print=True)

        sys.exit(0)

sys.exit(process.returncode)