diff --git a/ate/cli.py b/ate/cli.py index fd0482f2..4c37a546 100644 --- a/ate/cli.py +++ b/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() diff --git a/ate/locusts.py b/ate/locusts.py index b087e750..4d0cb6c0 100644 --- a/ate/locusts.py +++ b/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