Initial commit

.gitignore

@@ -0,0 +1,5 @@
+*.pyc
+__pycache__
+*.zip
+python-packages
+terraform.tfstate*

README.md

@@ -0,0 +1,140 @@
+# invokust
+
+A small wrapper for [locust](http://locust.io/) to enable invoking locust load tests from within Python. This gives more flexibility for automation such as QA/CI/CD tests and also makes it possible to run locust on [AWS Lambda](https://aws.amazon.com/lambda/).
+
+## Installation
+
+Install via pip:
+
+```
+pip install invokust
+```
+
+## Examples
+
+Running a load test using a locust file:
+
+```python
+import invokust
+
+settings = invokust.create_settings(
+    locustfile='locustfile_example.py',
+    host='http://example.com',
+    num_requests=10,
+    num_clients=1,
+    hatch_rate=1
+    )
+
+loadtest = invokust.LoadTest(settings)
+loadtest.run()
+loadtest.stats()
+'{"fail": {}, "locust_host": "http://example.com", "num_requests": 10, "success": {"/": {"num_requests": 10, "total_rps": 0.9611445710106717, "median_response_time": 110, "total_rpm": 57.6686742606403, "request_type": "GET", "min_response_time": 107, "max_response_time": 143}}}'
+```
+
+Running a load test without locust file:
+
+```python
+import invokust
+
+from locust import HttpLocust, TaskSet, task
+
+class Task(TaskSet):
+    @task()
+    def get_home_page(self):
+        '''
+        Gets /
+        '''
+        self.client.get("/")
+
+class WebsiteUser(HttpLocust):
+    task_set = Task
+
+settings = invokust.create_settings(
+    classes=[WebsiteUser],
+    host='http://example.com',
+    num_requests=10,
+    num_clients=1,
+    hatch_rate=1
+    )
+
+loadtest = invokust.LoadTest(settings)
+loadtest.run()
+loadtest.stats()
+'{"fail": {}, "locust_host": "http://example.com", "num_requests": 10, "success": {"/": {"num_requests": 10, "total_rps": 0.9806702027934636, "median_response_time": 110, "total_rpm": 58.84021216760782, "request_type": "GET", "min_response_time": 105, "max_response_time": 140}}}'
+```
+
+## Running on AWS Lambda
+
+<img src="http://d0.awsstatic.com/Graphics/lambda-icon-smallr1.png" alt="Lambda logo" height="100"><img src="http://locust.io/static/img/logo.png" alt="Locust logo" height="100">
+
+[AWS Lambda](https://aws.amazon.com/lambda/) is a great tool for load testing as it is very cheap and highly scalable.
+
+When calling the `create_settings` function, an argument can be passed to get the settings from environment variables. This allows the load test settings to be changed using the environment variables of the Lambda function:
+
+```python
+settings = invokust.create_settings(from_environment=True)
+```
+
+The environment variables are:
+
+  - LOCUST_LOCUSTFILE: Locust file to use for the load test
+  - LOCUST_CLASSES: Names of locust classes to use for the load test (instead of a locustfile). If more than one, separate with comma.
+  - LOCUST_HOST: The host to run the load test against
+  - LOCUST_NUM_REQUESTS: Total number of requests to make
+  - LOCUST_NUM_CLIENTS: Number of clients to simulate
+  - LOCUST_HATCH_RATE: Number of clients per second to start
+
+A time limit can also be placed on the load test to suit the Lambda execution time limit of 300 seconds. This will ensure you get the statistics even if the load test doesn't fully complete within 300 seconds:
+
+```python
+loadtest.run(timeout=298)
+```
+
+The load test statistics can be logged using standard python logging. The statistics are then recorded in [Cloudwatch Logs](https://eu-central-1.console.aws.amazon.com/cloudwatch/home?region=eu-central-1#logs) for summing, parsing or processing:
+
+```python
+logger.info(loadtest.stats())
+```
+
+A full example is in [aws_lambda_example.py](aws_lambda_example.py).
+
+### Creating a Lambda function
+
+The process for running a locust test on Lambda involves [creating a zip file](http://docs.aws.amazon.com/lambda/latest/dg/lambda-python-how-to-create-deployment-package.html) of the locust load test, creating a Lambda function and then triggering the function.
+
+Install invokust (and its dependencies) python packages locally:
+
+```
+pip install invokust --target=python-packages
+```
+
+Or if running on a Mac (python packages need to be compiled for 64 bit Linux) you can use docker:
+
+```
+docker run -it --volume=$PWD/python-packages:/python-packages python:2.7 bash -c "pip install invokust --target=/python-packages"
+```
+
+Create the zip file:
+
+```
+zip -q -r invokust_example.zip aws_lambda_example.py locustfile_example.py python-packages
+```
+
+Then create the Lambda function using [Terraform](https://www.terraform.io/) and the example [main.tf](main.tf) file:
+
+```
+terraform apply
+...
+```
+
+Finally invoke the function using the [AWS CLI](https://aws.amazon.com/cli/) (or use the [Lambda console](https://eu-central-1.console.aws.amazon.com/lambda/home?region=eu-central-1#/functions)):
+
+```
+aws lambda invoke --function-name invokust_example output.log
+{
+    "StatusCode": 200
+}
+
+cat output.log
+"{\"num_requests_fail\": 0, \"num_requests\": 10, \"success\": {\"/\": {\"num_requests\": 10, \"total_rps\": 1.0718836271241832, \"median_response_time\": 99, \"total_rpm\": 64.31301762745099, \"request_type\": \"GET\", \"min_response_time\": 95, \"max_response_time\": 100}}, \"locust_host\": \"http://example.com\", \"fail\": {}, \"num_requests_success\": 10}"
+```

TODO.md

@@ -0,0 +1,3 @@
+To do
+
+- `LOCUST_CLASSES` environment variable does not work

aws_lambda_example.py

@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+
+import sys
+sys.path.insert(0, "python-packages")
+
+import logging
+import invokust
+
+logging.basicConfig(level=logging.INFO)
+
+def lambda_handler(event=None, context=None):
+    try:
+        settings = invokust.create_settings(from_environment=True)
+        loadtest = invokust.LoadTest(settings)
+        loadtest.run(timeout=298)
+
+    except Exception as e:
+        logging.error("Locust exception {0}".format(repr(e)))
+
+    else:
+        return loadtest.stats()

invokust/__init__.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+
+from .settings import create_settings
+from .loadtest import LoadTest

invokust/loadtest.py

@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+
+import os
+import gevent
+import json
+import signal
+import logging
+
+from locust import runners, events
+
+
+logger = logging.getLogger(__name__)
+
+
+class TimeOutException(Exception): pass
+
+
+class LoadTest(object):
+    '''
+    The LoadTest runs the load test and returns statistics
+    '''
+    def __init__(self, settings):
+        self.settings = settings
+        signal.signal(signal.SIGALRM, self.sig_alarm_handler)
+        gevent.signal(signal.SIGTERM, self.sig_term_handler)
+
+    def stats(self):
+        '''
+        Returns the statistics from the load test in JSON
+        '''
+        statistics = {
+            'success': {},
+            'fail': {},
+            'num_requests': runners.locust_runner.stats.num_requests,
+            'num_requests_success': 0,
+            'num_requests_fail': 0,
+            'locust_host': runners.locust_runner.host
+        }
+
+        for name, value in runners.locust_runner.stats.entries.items():
+            statistics['success'][name[0]] = {
+                'request_type': name[1],
+                'num_requests': value.num_requests,
+                'min_response_time': value.min_response_time,
+                'median_response_time': value.median_response_time,
+                'max_response_time': value.max_response_time,
+                'total_rps': value.total_rps,
+                'total_rpm': value.total_rps * 60
+            }
+
+        for id, error in runners.locust_runner.errors.viewitems():
+            statistics['fail'][error.name] = error.to_dict()
+
+        statistics['num_requests_success'] = sum(
+            [statistics['success'][req]['num_requests'] for req in statistics['success']])
+        statistics['num_requests_fail'] = sum(
+            [statistics['fail'][req]['occurences'] for req in statistics['fail']])
+
+        return json.dumps(statistics)
+
+    def sig_term_handler(self, signum, frame):
+        logger.info("Received sigterm, exiting")
+        logger.info(self.stats())
+        sys.exit(0)
+
+    def sig_alarm_handler(self, signum, frame):
+        '''
+        This handler is used when a run time limit is set
+        '''
+        raise TimeOutException
+
+    def run(self, timeout=None):
+        '''
+        Run the load test. Optionally a timeout can be set to limit the run time
+        of the load test
+        '''
+        if timeout:
+            self.timeout = timeout
+            signal.alarm(timeout)
+        try:
+            logger.info("Starting Locust with and settings {0}".format(
+                vars(self.settings)))
+            runners.locust_runner = runners.LocalLocustRunner(self.settings.classes,
+                self.settings)
+            runners.locust_runner.start_hatching(wait=True)
+            runners.locust_runner.greenlet.join()
+            logger.info('Locust completed {0} requests with {1} errors'.format(
+                self.settings.num_requests,
+                len(runners.locust_runner.errors)))
+
+        except TimeOutException:
+            events.quitting.fire()
+            logger.info(self.stats())
+            logger.info("Run time limit reached: {0} seconds".format(self.timeout))
+
+        except Exception as e:
+            events.quitting.fire()
+            logger.error("Locust exception {0}".format(repr(e)))

invokust/settings.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+
+import os
+
+from locust.main import load_locustfile
+
+
+def create_settings(from_environment=False, locustfile=None,
+        classes=None, host=None, num_requests=None, num_clients=None,
+        hatch_rate=None):
+    '''
+    Returns a settings object to be used by a LocalLocustRunner.
+
+    Arguments
+
+        from_environment: get settings from environment variables
+        locustfile: locustfile to use for loadtest
+        classes: locust classes to use for load test
+        host: host for load testing
+        num_requests: number of requests to send
+        num_clients: number of clients to simulate in load test
+        hatch_rate: number of clients per second to start
+
+    If from_environment is set to True then this function will attempt to set
+    the attributes from environment variables. The environment variables are
+    named LOCUST_ + attribute name in upper case.
+    '''
+
+    settings = type('', (), {})()
+
+    settings.from_environment = from_environment
+    settings.locustfile = locustfile
+    settings.classes = classes
+    settings.host = host
+    settings.num_requests = num_requests
+    settings.num_clients = num_clients
+    settings.hatch_rate = hatch_rate
+
+    # Default settings that are not to be changed
+    settings.no_web = True
+    settings.master = False
+    settings.show_task_ratio_json = False
+    settings.list_commands = False
+    settings.loglevel = 'INFO'
+    settings.slave = False
+    settings.only_summary = True
+    settings.logfile = None
+    settings.show_task_ratio = False
+    settings.print_stats = False
+
+    if from_environment:
+        for attribute in ['locustfile', 'classes', 'host', 'num_requests', 'num_clients', 'hatch_rate']:
+            var_name = 'LOCUST_{0}'.format(attribute.upper())
+            var_value = os.environ.get(var_name)
+            if var_value and var_value.isdigit():
+                var_value = int(var_value)
+            setattr(settings, attribute, var_value)
+
+    if settings.locustfile is None and settings.classes is None:
+        raise Exception('One of locustfile or classes must be specified')
+
+    if settings.locustfile and settings.classes:
+        raise Exception('Only one of locustfile or classes can be specified')
+
+    if settings.locustfile:
+        docstring, classes = load_locustfile(settings.locustfile)
+        settings.classes = [classes[n] for n in classes]
+    else:
+        if isinstance(settings.classes, str):
+            settings.classes = settings.classes.split(',')
+            for idx, val in enumerate(settings.classes):
+                # This needs fixing
+                settings.classes[idx] = eval(val)
+
+    for attribute in ['classes', 'host', 'num_requests', 'num_clients', 'hatch_rate']:
+        val = getattr(settings, attribute, None)
+        if not val:
+            raise Exception('configuration error. LOCUST_{0} is not set'.format(attribute.upper()))
+
+    return settings

locustfile_example.py

@@ -0,0 +1,18 @@
+# -*- coding: utf-8 -*-
+
+import logging
+
+from locust import HttpLocust, TaskSet, task
+
+logger = logging.getLogger(__name__)
+
+class Task(TaskSet):
+    @task()
+    def get_home_page(self):
+        '''
+        Gets /
+        '''
+        self.client.get("/")
+
+class WebsiteUser(HttpLocust):
+    task_set = Task

main.tf

@@ -0,0 +1,22 @@
+# Example terraform configuration for a lambda function that uses invokust to run locust
+
+data "aws_caller_identity" "current" {}
+
+resource "aws_lambda_function" "invokust_example" {
+  filename      = "invokust_example.zip"
+  function_name = "invokust_example"
+  role          = "arn:aws:iam::${data.aws_caller_identity.current.account_id}:role/lambda_basic_execution"
+  handler       = "aws_lambda_example.lambda_handler"
+  runtime       = "python2.7"
+  timeout       = 300
+  description   = "A function that runs a locust load test"
+  environment {
+    variables = {
+      LOCUST_NUM_REQUESTS="10"
+      LOCUST_LOCUSTFILE="locustfile_example.py"
+      LOCUST_HOST="http://example.com"
+      LOCUST_HATCH_RATE="1"
+      LOCUST_NUM_CLIENTS="1"
+    }
+  }
+}