From d04207f889a6f0068a38fb4f16157f06f7431bc5 Mon Sep 17 00:00:00 2001
From: httprunner <httprunner>
Date: Wed, 31 Oct 2018 23:30:08 +0800
Subject: [PATCH] add --locustfile argument

---
 httprunner/cli.py | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/httprunner/cli.py b/httprunner/cli.py
index f0f1bda7..cfbdb0df 100644
--- a/httprunner/cli.py
+++ b/httprunner/cli.py
@@ -113,6 +113,10 @@ def main_locust():
     if len(sys.argv) == 1:
         sys.argv.extend(["-h"])
 
+    if sys.argv[1] in ["-h", "--help", "-V", "--version"]:
+        locusts.main()
+        sys.exit(0)
+
     # set logging level
     if "-L" in sys.argv:
         loglevel_index = sys.argv.index('-L') + 1
@@ -129,15 +133,18 @@ def main_locust():
 
     logger.setup_logger(loglevel)
 
-    if sys.argv[1] in ["-h", "--help", "-V", "--version"]:
-        locusts.main()
-        sys.exit(0)
-
+    # get testcase file path
     try:
-        testcase_index = sys.argv.index('-f') + 1
-        assert testcase_index < len(sys.argv)
-    except (ValueError, AssertionError):
-        logger.log_error("Testcase file is not specified, exit.")
+        if "-f" in sys.argv:
+            testcase_index = sys.argv.index('-f') + 1
+        elif "--locustfile" in sys.argv:
+            testcase_index = sys.argv.index('--locustfile') + 1
+        else:
+            testcase_index = None
+
+        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]