diff --git a/httprunner/ext/locusts/__init__.py b/httprunner/ext/locusts/__init__.py index 7c05d6f2..747fbfc1 100644 --- a/httprunner/ext/locusts/__init__.py +++ b/httprunner/ext/locusts/__init__.py @@ -29,6 +29,17 @@ def init_parser_locusts(subparsers): def main_locusts(args, extra_args): """ Performance test with locust: parse command line options and run commands. """ + try: + from locust.main import main + except ImportError: + msg = """ + Locust is not installed, install first and try again. + install with pip: + $ pip install locustio + """ + logger.error(msg) + sys.exit(1) + logger.info(f"HttpRunner version: {__version__}") sys.argv = ["locust", *extra_args] diff --git a/httprunner/ext/locusts/core.py b/httprunner/ext/locusts/core.py index c03049aa..51823c85 100644 --- a/httprunner/ext/locusts/core.py +++ b/httprunner/ext/locusts/core.py @@ -5,17 +5,6 @@ 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. @@ -59,7 +48,8 @@ def gen_locustfile(testcase_file_path): def start_locust_main(): logger.info(f"run command: {sys.argv}") - locust_main.main() + from locust.main import main + main() def start_master(sys_argv):