mirror of
https://github.com/httprunner/httprunner.git
synced 2026-05-12 02:21:29 +08:00
101 lines
2.6 KiB
Python
101 lines
2.6 KiB
Python
import io
|
|
import multiprocessing
|
|
import os
|
|
import sys
|
|
|
|
from loguru import logger
|
|
|
|
|
|
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):
|
|
logger.error("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
|
|
logger.error("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.py"
|
|
)
|
|
|
|
with io.open(template_path, encoding='utf-8') as template:
|
|
with io.open(locustfile_path, 'w', encoding='utf-8') as locustfile:
|
|
template_content = template.read()
|
|
template_content = template_content.replace("$TESTCASE_FILE", testcase_file_path)
|
|
locustfile.write(template_content)
|
|
|
|
return locustfile_path
|
|
|
|
|
|
def start_locust_main():
|
|
logger.info(f"run command: {sys.argv}")
|
|
from locust.main import main
|
|
main()
|
|
|
|
|
|
def start_master(sys_argv):
|
|
sys_argv.append("--master")
|
|
sys.argv = sys_argv
|
|
start_locust_main()
|
|
|
|
|
|
def start_slave(sys_argv):
|
|
if "--slave" not in sys_argv:
|
|
sys_argv.extend(["--slave"])
|
|
|
|
sys.argv = sys_argv
|
|
start_locust_main()
|
|
|
|
|
|
def init_slave_processes(slave_num):
|
|
""" init specified number of locust slave processes."""
|
|
processes = []
|
|
|
|
for _ in range(slave_num):
|
|
p_slave = multiprocessing.Process(target=start_slave, args=(sys.argv,))
|
|
p_slave.daemon = True
|
|
p_slave.start()
|
|
processes.append(p_slave)
|
|
|
|
return processes
|
|
|
|
|
|
def start_slaves(slave_num):
|
|
logger.info(f"Start {slave_num} locust slaves ...")
|
|
processes = init_slave_processes(slave_num)
|
|
[process.join() for process in processes]
|
|
|
|
|
|
def quick_run_locusts(slave_num):
|
|
""" quick start locust master and multiple slaves.
|
|
|
|
Args:
|
|
slave_num: locust slaves number
|
|
"""
|
|
logger.info(f"Start locust master with {slave_num} slaves ...")
|
|
|
|
processes = init_slave_processes(slave_num)
|
|
processes.append(
|
|
multiprocessing.Process(target=start_master, args=(sys.argv,))
|
|
)
|
|
[process.join() for process in processes]
|