From 1e64171fb8705b5ad7ce5e1c2241628b7b95f959 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Mon, 13 Apr 2020 00:33:11 +0800
Subject: [PATCH] fix: check if gevent installed when running locusts

---
 httprunner/cli.py                  | 14 ++++++++++++--
 httprunner/ext/locusts/__init__.py | 14 +++-----------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/httprunner/cli.py b/httprunner/cli.py
index a68cedc1..4aac1b87 100644
--- a/httprunner/cli.py
+++ b/httprunner/cli.py
@@ -4,8 +4,18 @@ import sys
 
 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()
+    try:
+        from gevent import monkey
+        monkey.patch_ssl()
+        from locust.main import main
+    except ImportError:
+        msg = """
+Locust is not installed, install first and try again.
+install with pip:
+$ pip install locustio
+"""
+        print(msg)
+        sys.exit(1)
 
 from loguru import logger
 
diff --git a/httprunner/ext/locusts/__init__.py b/httprunner/ext/locusts/__init__.py
index 747fbfc1..cdd6dfbb 100644
--- a/httprunner/ext/locusts/__init__.py
+++ b/httprunner/ext/locusts/__init__.py
@@ -29,17 +29,6 @@ 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]
 
@@ -86,5 +75,8 @@ def main_locusts(args, extra_args):
             start_master(sys.argv)
         elif args.slaves:
             start_slaves(args.slaves)
+        else:
+            quick_run_locusts(CPU_COUNT)
+
     except KeyboardInterrupt:
         manager.shutdown()