run locusts at full speed with master and several slaves, make the most use of all cpus

ate/__init__.py

@@ -1 +1 @@
-__version__ = '0.5.0'
+__version__ = '0.5.1'

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.

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)