"""Command-Line-Interface for wtg-seal.
This module provides a command line interface for wtg-seal.
Particularly, it manages some options and defines the main function.
"""
import argparse
import sys
from pathlib import Path
from typing import Dict, List
import locust.stats
from . import __version__ as wtgseal_version
from . import maker, utils
[docs]def add_default_args(parser: argparse.ArgumentParser) -> None:
"""Add default options and arguments to the CLI parser.
Parameters
----------
parser : {argparse.ArgumentParser}
CLI parser object
"""
parser.add_argument(
'docdef',
help='file with documents definitions',
metavar='DEFINITIONS')
parser.add_argument(
'docseq',
help='file with sequence of documents requests',
metavar='REQUESTS')
parser.add_argument(
'-l',
'--locust',
help='class name for Locust (default: WebsiteUser)',
default='WebsiteUser')
parser.add_argument(
'-t',
'--taskset',
help='class name for Taskset (default: UserBehaviour)',
default='UserBehaviour')
parser.add_argument(
'-w',
'--weight',
help='weight for Locust (default: 1)',
default=1,
type=int)
parser.add_argument(
'-s',
'--seed',
help='seed for wait time values generator (default: 1)',
default=1,
type=int,
metavar='SEED',
dest='wait_seed')
parser.add_argument(
'--csv-stats-interval',
help='the frequency (in seconds) with which stats data are written to '
'CSV file',
type=int,
metavar='INTERVAL',
dest='csv_stats_interval')
parser.add_argument(
'--group-by-doc',
help='group URLs together in Locust\'s statistics according to '
'document id',
action='store_true',
dest='group_by_doc')
parser.add_argument(
'-o',
'--output',
help='output file name (default: locustfile.py)',
metavar='OUTFILE',
default='locustfile.py')
parser.add_argument(
'-f',
'--force',
help='force overwriting an existing file',
action='store_true')
parser.add_argument(
'-V',
'--version',
action='version',
version=f'wtg-seal {wtgseal_version}')
[docs]def parse_args(args: List[str]) -> Dict:
"""Parse command line arguments.
Parameters
----------
args : {List[str]}
Command line arguments as list of strings.
Returns
-------
Dict
Command line arguments as dictionary.
"""
parser = argparse.ArgumentParser(
description='Create a locust file based on outputs from '
'SURGE for web traffic generation.',
epilog='Note: see the documentation for more information about SURGE '
'and expected format for input files.')
parser.set_defaults(command=run_wtgseal)
add_default_args(parser)
opts = vars(parser.parse_args(args))
return opts
[docs]def validate_file(path: Path) -> None:
"""Validate a given file path.
Validate if a given path exists and it is a file, throwing an
exception otherwise.
Parameters
----------
path : {Path}
The file represented as a Path object.
Raises
------
FileNotFoundError
If file does not exist or it is not a file.
"""
if not path.exists() or not path.is_file():
raise FileNotFoundError(f'{path} not found or it is not a file.')
[docs]def run_wtgseal(opts: Dict) -> None:
"""Create a locust file.
Call the python API to create a locust file, according to the
inputs provided.
Parameters
----------
opts : {Dict}
Command line arguments as a dictionary.
Raises
------
FileExistsError
If output file exists and option ``--force`` was not given.
"""
# setup input/output paths
docdef = Path(opts['docdef'])
docseq = Path(opts['docseq'])
outfile = Path(opts['output'])
# validate input files
validate_file(docdef)
validate_file(docseq)
# validate output file
if outfile.exists() and not opts['force']:
raise FileExistsError(f'{outfile} already exists.')
with docseq.open(mode='r') as fd:
freq = utils.count_requests(fd)
weights = utils.calc_weights(freq)
out = []
out.extend(maker.setup_header())
out.extend(maker.setup_blank_line())
out.extend(maker.setup_import())
if (opts['csv_stats_interval'] is not None
and opts['csv_stats_interval']
!= locust.stats.CSV_STATS_INTERVAL_SEC):
out.extend(maker.setup_csv_stats_interval(opts['csv_stats_interval']))
out.extend(maker.setup_blank_line())
out.extend(maker.setup_blank_line())
out.extend(maker.setup_taskset(opts['taskset']))
with docdef.open(mode='r') as fd:
docs = utils.parse_documents(fd)
for i, doc in enumerate(docs):
if opts['group_by_doc']:
out.extend(maker.setup_task(f'getdoc{i}', doc,
weight=weights[i], indlevel=1,
group_name=f'doc#{i}'))
else:
out.extend(maker.setup_task(f'getdoc{i}', doc,
weight=weights[i], indlevel=1))
out.extend(maker.setup_blank_line())
out.extend(maker.setup_blank_line())
out.extend(maker.setup_user(opts['locust'], opts['taskset'],
weight=opts['weight'], wait_seed=opts['wait_seed']))
maker.write_locust(outfile, out)
[docs]def main(args: List[str]):
"""Set main entry point for external applications.
Parameters
----------
args : {List[str]}
Command line arguments.
"""
opts = parse_args(args)
opts['command'](opts)
[docs]def run():
"""Entry point for console script."""
main(sys.argv[1:])
if __name__ == '__main__':
main(sys.argv[1:])