adjust code structure related to locust

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