run

#!/usr/bin/env python3
"""
Usage:
  run <base_url> [--persona=<persona>] [--duration=<length>] [--users=<cnt>] [--custom-output] [--debug] [--vrp-tomtom-probability <prob>] [--vrp-profiles <profiles>] [--points <points>] [--num-points <num_points>] [--vrp-max-profiles <vrp_max_profiles>] [--vrp-max-locations <vrp_max_locations>] [--master] [--expect-workers <clients_count>] [--worker] [--master-host <IP>] [--api-key <api_key>]

Options:
  -h --help                                Show this screen.
  <base_url>                               GraphHopper base URL, e.g. `https://example.com/api/1` or `https://localhost`.
  -d <length>, --duration <length>         Duration of the load-test, e.g. 60s, 20m [default: 1m].
  -u <cnt>, --users <cnt>                  Users count to simultaneously hit GraphHopper [default: 100].
  -p <persona> --persona=<persona>         The persona to use. Possible values: vrp, route, route_invalid, isochrone.
  --custom-output                          Use custom results output that does not spam with intermediate results.
  --debug                                  Switch debugging for error requests.
  --vrp-tomtom-probability <prob>          When doing VRP /optimize requests, what's the probability that an eligible requests (e.g. car) will use TomTom data [default: 0.2].
  --vrp-profiles <profiles>                A comma-separated list of allowed profiles, e.g. car,truck,small_truck. Possible values: car, bike, foot, hike, mtb, racingbike, scooter, truck, small_truck. [default: car]
  --points <points>                        A points query for matrix or routing, e.g. point=52.517429%2C13.389083&point=52.522018%2C13.386669.
  --num-points <num_points>                The number of points for matrix and routing queries. This setting is mutually exclusive with the --points option.
  --vrp-max-profiles <vrp_max_profiles>    The max of different VRP vehicle profiles to use. [default: 1]
  --vrp-max-locations <vrp_max_locations>  The max locations for the services. [default: 10]
  --master                                 Should this run be a master in a master/worker configuration with multiple clients?
  --expect-workers <clients_count>         If master is set, how many clients do we expect?
  --worker                                 Should this run be a worker in a master/worker configuration?
  --master-host <IP>                       If worker is set, where is the master running?
  --api-key <api_key>                      Use a custom Dirigent API key

"""
from docopt import docopt
from utils import color
import main
import os
import run_as_cli
import run_as_lib
import signal
import sys
import time


CURRENT_VERSION = "v1.1.0"


def get_settings():
    # set arguments
    args = docopt(__doc__)
    settings = {}
    settings["base_url"] = (
        args["<base_url>"] if "http" in args["<base_url>"] else f"https://{args['<base_url>']}"
    ).strip("/")
    settings["users"] = args["--users"]
    settings["duration"] = args["--duration"]
    settings["custom-output"] = args["--custom-output"]
    settings["debug"] = args["--debug"]

    # set persona
    persona_mapping = {
        "vrp": "PersonaVRP",
        "matrix": "PersonaMatrix",
        "route": "PersonaRoute",
        "route_invalid": "PersonaRouteInvalid",
        "isochrone": "PersonaIsochrone",
    }
    persona = args["--persona"]
    if persona not in persona_mapping:
        available = ", ".join(list(persona_mapping.keys()))
        sys.exit(f"Persona {persona} is not supported, you can use any of [{available}].")
    settings["persona_arg"] = persona_mapping[persona]

    # set custom class settings
    # there's no way to pass params to Locust classes, so an environment will have to do
    vrp_tomtom_probability = float(args["--vrp-tomtom-probability"])
    if vrp_tomtom_probability > 1.0 or vrp_tomtom_probability < 0.0:
        print("The --vrp-tomtom-probability parameter must be within 0.0 and 1.0.")
        sys.exit(1)
    os.environ["VRP_TOMTOM_PROBABILITY"] = str(vrp_tomtom_probability)
    os.environ["VRP_PROFILES"] = args["--vrp-profiles"]
    if args["--points"]:
        os.environ["QUERY_POINTS"] = args["--points"]
    if args["--num-points"]:
        os.environ["QUERY_POINTS_NUM"] = args["--num-points"]
    if args["--api-key"]:
        os.environ["API_KEY"] = args["--api-key"]
    os.environ["VRP_MAX_PROFILES"] = args["--vrp-max-profiles"]
    os.environ["VRP_MAX_LOCATIONS"] = args["--vrp-max-locations"]
    os.environ["DEBUG"] = "yes" if args["--debug"] else "no"

    # master/worker
    settings["master"] = args["--master"]
    settings["expect-workers"] = args["--expect-workers"]
    settings["worker"] = args["--worker"]
    settings["master-host"] = args["--master-host"]

    os.environ["PYTHONWARNINGS"] = "ignore:Unverified HTTPS request"

    return settings


def run_locust(settings):
    if settings["custom-output"]:
        run_as_lib.run(settings)
    else:
        run_as_cli.run(settings)


def print_info(settings):
    print(color.green_bg("Running Locust."))
    print("")
    print("Base URL: {}".format(color.green_bg(settings["base_url"])))
    print("Users:    {}".format(color.green_bg(settings["users"])))
    print("Duration: {}".format(color.green_bg(settings["duration"])))
    print("Persona:  {}".format(color.green_bg(settings["persona_arg"])))


def setup_ctrlc_signal_handler():
    def ctrlc_signal_handler(signal, frame):
        sys.exit(0)
    signal.signal(signal.SIGINT, ctrlc_signal_handler)


if __name__ == "__main__":
    setup_ctrlc_signal_handler()
    print("Load test for GraphHopper services.")
    print(f"Version: {CURRENT_VERSION}.")
    print("")
    settings = get_settings()
    print_info(settings)
    run_locust(settings)