diff --git a/ate/__init__.py b/ate/__init__.py index 9bdd4d27..08d79c0e 100644 --- a/ate/__init__.py +++ b/ate/__init__.py @@ -1 +1 @@ -__version__ = '0.5.0' \ No newline at end of file +__version__ = '0.5.1' \ No newline at end of file diff --git a/ate/cli.py b/ate/cli.py index 4f30f7e6..fd0482f2 100644 --- a/ate/cli.py +++ b/ate/cli.py @@ -4,11 +4,12 @@ import logging import os import sys from collections import OrderedDict -import PyUnitReport from ate import __version__ from ate.task import create_task +import PyUnitReport + def main_ate(): """ API test: parse command line options and run commands. @@ -88,7 +89,7 @@ def main_locust(): """ Performance test with locust: parse command line options and run commands. """ try: - from locust.main import main + from ate.locusts import main, run_locusts_at_full_speed except ImportError: print("Locust is not installed, exit.") exit(1) @@ -110,7 +111,11 @@ def main_locust(): testcase_file_path = sys.argv[testcase_index] sys.argv[testcase_index] = parse_locustfile(testcase_file_path) - main() + + if "--full-speed" in sys.argv: + run_locusts_at_full_speed(sys.argv) + else: + main() def parse_locustfile(file_path): """ parse testcase file and return locustfile path. diff --git a/ate/locusts.py b/ate/locusts.py new file mode 100644 index 00000000..b087e750 --- /dev/null +++ b/ate/locusts.py @@ -0,0 +1,31 @@ +import multiprocessing +import sys + +from locust.main import main + + +def start_master(sys_argv): + sys_argv.append("--master") + sys.argv = sys_argv + main() + +def start_slave(sys_argv): + sys_argv.extend(["--slave"]) + sys.argv = sys_argv + main() + +def run_locusts_at_full_speed(sys_argv): + sys_argv.pop(sys_argv.index("--full-speed")) + slaves_num = multiprocessing.cpu_count() + + processes = [] + for _ in range(slaves_num): + p_slave = multiprocessing.Process(target=start_slave, args=(sys_argv,)) + p_slave.daemon = True + p_slave.start() + processes.append(p_slave) + + try: + start_master(sys_argv) + except KeyboardInterrupt: + sys.exit(0)