change: make locusts as hrun sub-command, usage: hrun locusts -h

docs/CHANGELOG.md

@@ -8,6 +8,7 @@
 - remove compatibility with testcase/testsuite format v1
 - make `startproject` as hrun sub-command, usage: `hrun startproject <project_name>`
 - make `har2case` as hrun sub-command, usage: `hrun har2case -h`
+- make `locusts` as hrun sub-command, usage: `hrun locusts -h`
 
 ## 3.0.1 (2020-03-24)
 

httprunner/cli.py

@@ -4,10 +4,16 @@ import sys
 
 from loguru import logger
 
+if len(sys.argv) >= 2 and sys.argv[1] == "locusts":
+    # monkey patch ssl at beginning to avoid RecursionError when running locust.
+    from gevent import monkey
+    monkey.patch_ssl()
+
 from httprunner import __description__, __version__
 from httprunner.api import HttpRunner
 from httprunner.ext.har2case import init_har2case_parser, main_har2case
 from httprunner.ext.scaffold import init_parser_scaffold, main_scaffold
+from httprunner.ext.locusts import init_parser_locusts, main_locusts
 from httprunner.report import gen_html_report
 
 
@@ -85,8 +91,13 @@ def main():
     sub_parser_run = init_parser_run(subparsers)
     sub_parser_scaffold = init_parser_scaffold(subparsers)
     sub_parser_har2case = init_har2case_parser(subparsers)
+    sub_parser_locusts = init_parser_locusts(subparsers)
 
-    args = parser.parse_args()
+    extra_args = []
+    if len(sys.argv) >= 2 and sys.argv[1] == "locusts":
+        args, extra_args = parser.parse_known_args()
+    else:
+        args = parser.parse_args()
 
     if args.version:
         print(f"{__version__}")
@@ -121,6 +132,14 @@ def main():
 
         main_har2case(args)
 
+    elif sys.argv[1] == "locusts":
+        # hrun locusts
+        if len(sys.argv) == 2:
+            sub_parser_locusts.print_help()
+            sys.exit(0)
+
+        main_locusts(args, extra_args)
+
 
 if __name__ == '__main__':
     main()

httprunner/ext/locusts/__init__.py

@@ -0,0 +1,79 @@
+import multiprocessing
+import sys
+
+from loguru import logger
+
+from httprunner import __version__
+from httprunner.ext.locusts.core import start_locust_main, parse_locustfile, quick_run_locusts, start_master, \
+    start_slaves
+
+CPU_COUNT = multiprocessing.cpu_count()
+
+
+def init_parser_locusts(subparsers):
+    sub_parser_locusts = subparsers.add_parser(
+        "locusts", help="Run load test with locust.")
+    sub_parser_locusts.add_argument(
+        '--locust-help', action='store_true', default=False,
+        help="Show locust help.")
+    sub_parser_locusts.add_argument(
+        "--master", action='store_true', default=False, help="Start locust master.")
+    sub_parser_locusts.add_argument(
+        "--slaves", type=int, help="Specify locust slave number.")
+    sub_parser_locusts.add_argument(
+        "--quickstart", action='store_true', default=False,
+        help=f"Start locust master with {CPU_COUNT} slaves.")
+    return sub_parser_locusts
+
+
+def main_locusts(args, extra_args):
+    """ Performance test with locust: parse command line options and run commands.
+    """
+    logger.info(f"HttpRunner version: {__version__}")
+    sys.argv = ["locust", *extra_args]
+
+    if args.locust_help:
+        sys.argv = ["locust", "-h"]
+        start_locust_main()
+
+    def get_arg_index(*target_args):
+        for arg in target_args:
+            if arg not in sys.argv:
+                continue
+
+            return sys.argv.index(arg) + 1
+
+        return None
+
+    # set logging level
+    loglevel_index = get_arg_index("-L", "--loglevel")
+    if loglevel_index and loglevel_index < len(sys.argv):
+        loglevel = sys.argv[loglevel_index]
+        loglevel = loglevel.upper()
+    else:
+        # default
+        loglevel = "INFO"
+
+    logger.remove()
+    logger.add(sys.stdout, level=loglevel)
+
+    # convert httprunner yaml/json case to locustfile.py
+    try:
+        testcase_index = get_arg_index("-f", "--locustfile")
+        assert testcase_index and testcase_index < len(sys.argv)
+        testcase_file_path = sys.argv[testcase_index]
+        sys.argv[testcase_index] = parse_locustfile(testcase_file_path)
+    except AssertionError:
+        print("Testcase file is not specified, exit.")
+        sys.exit(1)
+
+    manager = multiprocessing.Manager()
+    try:
+        if args.quickstart:
+            quick_run_locusts(CPU_COUNT)
+        elif args.master:
+            start_master(sys.argv)
+        elif args.slaves:
+            start_slaves(args.slaves)
+    except KeyboardInterrupt:
+        manager.shutdown()

httprunner/ext/locusts/__main__.py

@@ -1,4 +0,0 @@
-from httprunner.ext.locusts.cli import main
-
-if __name__ == "__main__":
-    main()

httprunner/ext/locusts/cli.py

@@ -1,174 +0,0 @@
-try:
-    # monkey patch ssl at beginning to avoid RecursionError when running locust.
-    from gevent import monkey
-    monkey.patch_ssl()
-    from locust import main as locust_main
-except ImportError:
-    msg = """
-Locust is not installed, install first and try again.
-install with pip:
-$ pip install locustio
-"""
-    print(msg)
-    import sys
-    sys.exit(0)
-
-import io
-import multiprocessing
-import os
-import sys
-
-from loguru import logger
-
-from httprunner import __version__
-
-
-def parse_locustfile(file_path):
-    """ parse testcase file and return locustfile path.
-        if file_path is a Python file, assume it is a locustfile
-        if file_path is a YAML/JSON file, convert it to locustfile
-    """
-    if not os.path.isfile(file_path):
-        logger.error("file path invalid, exit.")
-        sys.exit(1)
-
-    file_suffix = os.path.splitext(file_path)[1]
-    if file_suffix == ".py":
-        locustfile_path = file_path
-    elif file_suffix in ['.yaml', '.yml', '.json']:
-        locustfile_path = gen_locustfile(file_path)
-    else:
-        # '' or other suffix
-        logger.error("file type should be YAML/JSON/Python, exit.")
-        sys.exit(1)
-
-    return locustfile_path
-
-
-def gen_locustfile(testcase_file_path):
-    """ generate locustfile from template.
-    """
-    locustfile_path = 'locustfile.py'
-    template_path = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)),
-        "locustfile_template.py"
-    )
-
-    with io.open(template_path, encoding='utf-8') as template:
-        with io.open(locustfile_path, 'w', encoding='utf-8') as locustfile:
-            template_content = template.read()
-            template_content = template_content.replace("$TESTCASE_FILE", testcase_file_path)
-            locustfile.write(template_content)
-
-    return locustfile_path
-
-
-def start_locust_main():
-    locust_main.main()
-
-
-def start_master(sys_argv):
-    sys_argv.append("--master")
-    sys.argv = sys_argv
-    start_locust_main()
-
-
-def start_slave(sys_argv):
-    if "--slave" not in sys_argv:
-        sys_argv.extend(["--slave"])
-
-    sys.argv = sys_argv
-    start_locust_main()
-
-
-def run_locusts_with_processes(sys_argv, processes_count):
-    processes = []
-    manager = multiprocessing.Manager()
-
-    for _ in range(processes_count):
-        p_slave = multiprocessing.Process(target=start_slave, args=(sys_argv,))
-        p_slave.daemon = True
-        p_slave.start()
-        processes.append(p_slave)
-
-    try:
-        if "--slave" in sys_argv:
-            [process.join() for process in processes]
-        else:
-            start_master(sys_argv)
-    except KeyboardInterrupt:
-        manager.shutdown()
-
-
-def main():
-    """ Performance test with locust: parse command line options and run commands.
-    """
-    print(f"HttpRunner version: {__version__}")
-    sys.argv[0] = 'locust'
-    if len(sys.argv) == 1:
-        sys.argv.extend(["-h"])
-
-    if sys.argv[1] in ["-h", "--help", "-V", "--version"]:
-        start_locust_main()
-
-    def get_arg_index(*target_args):
-        for arg in target_args:
-            if arg not in sys.argv:
-                continue
-
-            return sys.argv.index(arg) + 1
-
-        return None
-
-    # set logging level
-    loglevel_index = get_arg_index("-L", "--loglevel")
-    if loglevel_index and loglevel_index < len(sys.argv):
-        loglevel = sys.argv[loglevel_index]
-        loglevel = loglevel.upper()
-    else:
-        # default
-        loglevel = "WARNING"
-
-    logger.remove()
-    logger.add(sys.stdout, level=loglevel)
-
-    # get testcase file path
-    try:
-        testcase_index = get_arg_index("-f", "--locustfile")
-        assert testcase_index and testcase_index < len(sys.argv)
-    except AssertionError:
-        print("Testcase file is not specified, exit.")
-        sys.exit(1)
-
-    testcase_file_path = sys.argv[testcase_index]
-    sys.argv[testcase_index] = parse_locustfile(testcase_file_path)
-
-    if "--processes" in sys.argv:
-        """ locusts -f locustfile.py --processes 4
-        """
-        if "--no-web" in sys.argv:
-            logger.error("conflict parameter args: --processes & --no-web. \nexit.")
-            sys.exit(1)
-
-        processes_index = sys.argv.index('--processes')
-        processes_count_index = processes_index + 1
-        if processes_count_index >= len(sys.argv):
-            """ do not specify processes count explicitly
-                locusts -f locustfile.py --processes
-            """
-            processes_count = multiprocessing.cpu_count()
-            logger.warning(f"processes count not specified, use {processes_count} by default.")
-        else:
-            try:
-                """ locusts -f locustfile.py --processes 4 """
-                processes_count = int(sys.argv[processes_count_index])
-                sys.argv.pop(processes_count_index)
-            except ValueError:
-                """ locusts -f locustfile.py --processes -P 8888 """
-                processes_count = multiprocessing.cpu_count()
-                logger.warning(f"processes count not specified, use {processes_count} by default.")
-
-        sys.argv.pop(processes_index)
-        run_locusts_with_processes(sys.argv, processes_count)
-    else:
-        start_locust_main()

httprunner/ext/locusts/core.py

@@ -0,0 +1,110 @@
+import io
+import multiprocessing
+import os
+import sys
+
+from loguru import logger
+
+try:
+    from locust import main as locust_main
+except ImportError:
+    msg = """
+Locust is not installed, install first and try again.
+install with pip:
+$ pip install locustio
+"""
+    logger.error(msg)
+    sys.exit(0)
+
+
+def parse_locustfile(file_path):
+    """ parse testcase file and return locustfile path.
+        if file_path is a Python file, assume it is a locustfile
+        if file_path is a YAML/JSON file, convert it to locustfile
+    """
+    if not os.path.isfile(file_path):
+        logger.error("file path invalid, exit.")
+        sys.exit(1)
+
+    file_suffix = os.path.splitext(file_path)[1]
+    if file_suffix == ".py":
+        locustfile_path = file_path
+    elif file_suffix in ['.yaml', '.yml', '.json']:
+        locustfile_path = gen_locustfile(file_path)
+    else:
+        # '' or other suffix
+        logger.error("file type should be YAML/JSON/Python, exit.")
+        sys.exit(1)
+
+    return locustfile_path
+
+
+def gen_locustfile(testcase_file_path):
+    """ generate locustfile from template.
+    """
+    locustfile_path = 'locustfile.py'
+    template_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "locustfile_template.py"
+    )
+
+    with io.open(template_path, encoding='utf-8') as template:
+        with io.open(locustfile_path, 'w', encoding='utf-8') as locustfile:
+            template_content = template.read()
+            template_content = template_content.replace("$TESTCASE_FILE", testcase_file_path)
+            locustfile.write(template_content)
+
+    return locustfile_path
+
+
+def start_locust_main():
+    logger.info(f"run command: {sys.argv}")
+    locust_main.main()
+
+
+def start_master(sys_argv):
+    sys_argv.append("--master")
+    sys.argv = sys_argv
+    start_locust_main()
+
+
+def start_slave(sys_argv):
+    if "--slave" not in sys_argv:
+        sys_argv.extend(["--slave"])
+
+    sys.argv = sys_argv
+    start_locust_main()
+
+
+def init_slave_processes(slave_num):
+    """ init specified number of locust slave processes."""
+    processes = []
+
+    for _ in range(slave_num):
+        p_slave = multiprocessing.Process(target=start_slave, args=(sys.argv,))
+        p_slave.daemon = True
+        p_slave.start()
+        processes.append(p_slave)
+
+    return processes
+
+
+def start_slaves(slave_num):
+    logger.info(f"Start {slave_num} locust slaves ...")
+    processes = init_slave_processes(slave_num)
+    [process.join() for process in processes]
+
+
+def quick_run_locusts(slave_num):
+    """ quick start locust master and multiple slaves.
+
+    Args:
+        slave_num: locust slaves number
+    """
+    logger.info(f"Start locust master with {slave_num} slaves ...")
+
+    processes = init_slave_processes(slave_num)
+    processes.append(
+        multiprocessing.Process(target=start_master, args=(sys.argv,))
+    )
+    [process.join() for process in processes]

httprunner/ext/locusts/utils_test.py

@@ -8,7 +8,7 @@ class TestLocust(unittest.TestCase):
 
     def test_prepare_locust_tests(self):
         path = os.path.join(
-            os.getcwd(), 'tests/locust_tests/demo_locusts.yml')
+            os.path.dirname(__file__), "data", "demo_locusts.yml")
         locust_tests = prepare_locust_tests(path)
         self.assertEqual(len(locust_tests), 2 + 3)
         name_list = [

pyproject.toml

@@ -48,7 +48,6 @@ fastapi = "^0.49.0"
 hrun = "httprunner.cli:main"
 ate = "httprunner.cli:main"
 httprunner = "httprunner.cli:main"
-locusts = "httprunner.ext.locusts.cli:main"
 
 [build-system]
 requires = ["poetry>=1.0.0"]