From b9392d7cce0e57048726441e7399faedf0c12d21 Mon Sep 17 00:00:00 2001 From: debugtalk Date: Sun, 12 Apr 2020 23:29:00 +0800 Subject: [PATCH] change: make `locusts` as hrun sub-command, usage: `hrun locusts -h` --- docs/CHANGELOG.md | 1 + httprunner/cli.py | 21 ++- httprunner/ext/locusts/__init__.py | 79 ++++++++ httprunner/ext/locusts/__main__.py | 4 - httprunner/ext/locusts/cli.py | 174 ------------------ httprunner/ext/locusts/core.py | 110 +++++++++++ .../ext/locusts/data}/demo_locusts.yml | 0 .../ext/locusts/utils_test.py | 2 +- pyproject.toml | 1 - tests/test_extension/__init__.py | 0 10 files changed, 211 insertions(+), 181 deletions(-) delete mode 100644 httprunner/ext/locusts/__main__.py delete mode 100644 httprunner/ext/locusts/cli.py create mode 100644 httprunner/ext/locusts/core.py rename {tests/locust_tests => httprunner/ext/locusts/data}/demo_locusts.yml (100%) rename tests/test_extension/test_locusts.py => httprunner/ext/locusts/utils_test.py (89%) delete mode 100644 tests/test_extension/__init__.py diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 694c0d28..981c1f59 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -8,6 +8,7 @@ - remove compatibility with testcase/testsuite format v1 - make `startproject` as hrun sub-command, usage: `hrun startproject ` - 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) diff --git a/httprunner/cli.py b/httprunner/cli.py index 85dd5ea2..f2fcdece 100644 --- a/httprunner/cli.py +++ b/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() diff --git a/httprunner/ext/locusts/__init__.py b/httprunner/ext/locusts/__init__.py index e69de29b..7c05d6f2 100644 --- a/httprunner/ext/locusts/__init__.py +++ b/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() diff --git a/httprunner/ext/locusts/__main__.py b/httprunner/ext/locusts/__main__.py deleted file mode 100644 index bc8d706f..00000000 --- a/httprunner/ext/locusts/__main__.py +++ /dev/null @@ -1,4 +0,0 @@ -from httprunner.ext.locusts.cli import main - -if __name__ == "__main__": - main() diff --git a/httprunner/ext/locusts/cli.py b/httprunner/ext/locusts/cli.py deleted file mode 100644 index 22e80c7c..00000000 --- a/httprunner/ext/locusts/cli.py +++ /dev/null @@ -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() diff --git a/httprunner/ext/locusts/core.py b/httprunner/ext/locusts/core.py new file mode 100644 index 00000000..c03049aa --- /dev/null +++ b/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] diff --git a/tests/locust_tests/demo_locusts.yml b/httprunner/ext/locusts/data/demo_locusts.yml similarity index 100% rename from tests/locust_tests/demo_locusts.yml rename to httprunner/ext/locusts/data/demo_locusts.yml diff --git a/tests/test_extension/test_locusts.py b/httprunner/ext/locusts/utils_test.py similarity index 89% rename from tests/test_extension/test_locusts.py rename to httprunner/ext/locusts/utils_test.py index 59d957ff..c76b0b52 100644 --- a/tests/test_extension/test_locusts.py +++ b/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 = [ diff --git a/pyproject.toml b/pyproject.toml index 943d0b91..3de0634c 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/test_extension/__init__.py b/tests/test_extension/__init__.py deleted file mode 100644 index e69de29b..00000000