From 389e915a09e006fb00cb38ffdd9ba133ff6a731f Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Thu, 11 Jan 2018 23:03:13 +0800
Subject: [PATCH] replace --cpu-cores with --full-speed: now support specify
 slaves number

---
 docs/load-test.md      |  4 ++--
 httprunner/__init__.py |  2 +-
 httprunner/cli.py      | 40 ++++++++++++++++++++++++++++++++--------
 httprunner/locusts.py  |  7 ++-----
 4 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/docs/load-test.md b/docs/load-test.md
index 2ac5e8b5..89e135a5 100644
--- a/docs/load-test.md
+++ b/docs/load-test.md
@@ -21,10 +21,10 @@ $ locusts -f examples/first-testcase.yml
 
 In this case, you can reuse all features of [`Locust`][Locust].
 
-That’s not all about it. With the argument `--full-speed`, you can even start locust with master and several slaves (default to cpu cores number) at one time, which means you can leverage all cpus of your machine.
+That’s not all about it. With the argument `--cpu-cores`, you can even start locust with master and specified number of slaves (default to cpu cores number) at one time, which means you can leverage all cpus of your machine.
 
 ```bash
-$ locusts -f examples/first-testcase.yml --full-speed
+$ locusts -f examples/first-testcase.yml --cpu-cores 4
 [2017-08-26 23:51:47,071] bogon/INFO/locust.main: Starting web monitor at *:8089
 [2017-08-26 23:51:47,075] bogon/INFO/locust.main: Starting Locust 0.8a2
 [2017-08-26 23:51:47,078] bogon/INFO/locust.main: Starting Locust 0.8a2
diff --git a/httprunner/__init__.py b/httprunner/__init__.py
index dd499f47..ff27c5be 100644
--- a/httprunner/__init__.py
+++ b/httprunner/__init__.py
@@ -1 +1 @@
-__version__ = '0.8.6'
\ No newline at end of file
+__version__ = '0.8.7'
\ No newline at end of file
diff --git a/httprunner/cli.py b/httprunner/cli.py
index 935d1f1d..ac3a8421 100644
--- a/httprunner/cli.py
+++ b/httprunner/cli.py
@@ -1,16 +1,16 @@
 import argparse
 import logging
+import multiprocessing
 import os
 import sys
 from collections import OrderedDict
 
-from pyunitreport import __version__ as pyu_version
-from pyunitreport import HTMLTestRunner
-
 from httprunner import __version__ as ate_version
 from httprunner import exception
 from httprunner.task import TaskSuite
-from httprunner.utils import create_scaffold
+from httprunner.utils import create_scaffold, string_type
+from pyunitreport import __version__ as pyu_version
+from pyunitreport import HTMLTestRunner
 
 
 def main_ate():
@@ -98,6 +98,8 @@ def main_ate():
 def main_locust():
     """ Performance test with locust: parse command line options and run commands.
     """
+    logging.basicConfig(level="INFO")
+
     try:
         from httprunner import locusts
     except ImportError:
@@ -124,12 +126,34 @@ def main_locust():
     testcase_file_path = sys.argv[testcase_index]
     sys.argv[testcase_index] = locusts.parse_locustfile(testcase_file_path)
 
-    if "--full-speed" in sys.argv:
-
+    if "--cpu-cores" in sys.argv:
+        """ locusts -f locustfile.py --cpu-cores 4
+        """
         if "--no-web" in sys.argv:
-            logging.warning("conflict parameter args: --full-speed --no-web. \nexit.")
+            logging.error("conflict parameter args: --cpu-cores & --no-web. \nexit.")
             sys.exit(1)
 
-        locusts.run_locusts_at_full_speed(sys.argv)
+        cpu_cores_index = sys.argv.index('--cpu-cores')
+
+        cpu_cores_num_index = cpu_cores_index + 1
+
+        if cpu_cores_num_index >= len(sys.argv):
+            """ do not specify cpu cores explicitly
+                locusts -f locustfile.py --cpu-cores
+            """
+            cpu_cores_num_value = multiprocessing.cpu_count()
+            logging.warning("cpu cores number not specified, use {} by default.".format(cpu_cores_num_value))
+        else:
+            try:
+                """ locusts -f locustfile.py --cpu-cores 4 """
+                cpu_cores_num_value = int(sys.argv[cpu_cores_num_index])
+                sys.argv.pop(cpu_cores_num_index)
+            except ValueError:
+                """ locusts -f locustfile.py --cpu-cores -P 8888 """
+                cpu_cores_num_value = multiprocessing.cpu_count()
+                logging.warning("cpu cores number not specified, use {} by default.".format(cpu_cores_num_value))
+
+        sys.argv.pop(cpu_cores_index)
+        locusts.run_locusts_on_cpu_cores(sys.argv, cpu_cores_num_value)
     else:
         locusts.main()
diff --git a/httprunner/locusts.py b/httprunner/locusts.py
index 5611efb2..73de4e89 100644
--- a/httprunner/locusts.py
+++ b/httprunner/locusts.py
@@ -58,12 +58,9 @@ def start_slave(sys_argv):
     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()
-
+def run_locusts_on_cpu_cores(sys_argv, cpu_cores_num_value):
     processes = []
-    for _ in range(slaves_num):
+    for _ in range(cpu_cores_num_value):
         p_slave = multiprocessing.Process(target=start_slave, args=(sys_argv,))
         p_slave.daemon = True
         p_slave.start()