adjust code structure related to locust

ate/cli.py

@@ -1,5 +1,4 @@
 import argparse
-import codecs
 import logging
 import os
 import sys
@@ -89,7 +88,7 @@ def main_locust():
     """ Performance test with locust: parse command line options and run commands.
     """
     try:
-        from ate.locusts import main, run_locusts_at_full_speed
+        from ate import locusts
     except ImportError:
         print("Locust is not installed, exit.")
         exit(1)
@@ -99,7 +98,7 @@ def main_locust():
         sys.argv.extend(["-h"])
 
     if sys.argv[1] in ["-h", "--help", "-V", "--version"]:
-        main()
+        locusts.main()
         sys.exit(0)
 
     try:
@@ -110,47 +109,9 @@ def main_locust():
         sys.exit(1)
 
     testcase_file_path = sys.argv[testcase_index]
-    sys.argv[testcase_index] = parse_locustfile(testcase_file_path)
+    sys.argv[testcase_index] = locusts.parse_locustfile(testcase_file_path)
 
     if "--full-speed" in sys.argv:
-        run_locusts_at_full_speed(sys.argv)
+        locusts.run_locusts_at_full_speed(sys.argv)
     else:
-        main()
-
-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):
-        print("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
-        print("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'
-    )
-    with codecs.open(template_path, encoding='utf-8') as template:
-        with codecs.open(locustfile_path, 'w', encoding='utf-8') as locustfile:
-            template_content = template.read()
-            template_content = template_content.replace("$HOST", "https://skypixel.com")
-            template_content = template_content.replace("$TESTCASE_FILE", testcase_file_path)
-            locustfile.write(template_content)
-
-    return locustfile_path
+        locusts.main()

ate/locusts.py

@@ -1,9 +1,49 @@
+import codecs
 import multiprocessing
+import os
 import sys
 
 from locust.main import main
 
 
+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):
+        print("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
+        print("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'
+    )
+    with codecs.open(template_path, encoding='utf-8') as template:
+        with codecs.open(locustfile_path, 'w', encoding='utf-8') as locustfile:
+            template_content = template.read()
+            template_content = template_content.replace("$HOST", "https://skypixel.com")
+            template_content = template_content.replace("$TESTCASE_FILE", testcase_file_path)
+            locustfile.write(template_content)
+
+    return locustfile_path
+
 def start_master(sys_argv):
     sys_argv.append("--master")
     sys.argv = sys_argv