def main() -> None:
language_choices = sorted(scenario_config.LANGUAGES.keys())
argp = argparse.ArgumentParser(description='Generates load test configs.')
argp.add_argument('-l',
'--language',
choices=language_choices,
required=True,
help='Language to benchmark.')
argp.add_argument('-t',
'--template',
type=str,
required=True,
help='LoadTest configuration yaml file template.')
argp.add_argument('-s',
'--substitutions',
action='extend',
nargs='+',
default=[],
type=str,
help='Template substitutions, in the form key=value.')
argp.add_argument('-p',
'--prefix',
default='',
type=str,
help='Test name prefix.')
argp.add_argument('-u',
'--uniquifiers',
action='extend',
nargs='+',
default=[],
type=str,
help='One or more strings to make the test name unique.')
argp.add_argument(
'-d',
nargs='?',
const=True,
default=False,
type=bool,
help='Use creation date and time as an addditional uniquifier.')
argp.add_argument('-a',
'--annotations',
action='extend',
nargs='+',
default=[],
type=str,
help='Test annotations, in the form key=value.')
argp.add_argument('-r',
'--regex',
default='.*',
type=str,
help='Regex to select scenarios to run.')
argp.add_argument(
'--category',
choices=['all', 'inproc', 'scalable', 'smoketest', 'sweep'],
default='all',
help='Select a category of tests to run.')
argp.add_argument(
'--client_language',
choices=language_choices,
help='Select only scenarios with a specified client language.')
argp.add_argument(
'--server_language',
choices=language_choices,
help='Select only scenarios with a specified server language.')
argp.add_argument('--runs_per_test',
default=1,
type=int,
help='Number of copies to generate for each test.')
argp.add_argument('-o',
'--output',
type=str,
help='Output file name. Output to stdout if not set.')
args = argp.parse_args()
substitutions = parse_key_value_args(args.substitutions)
with open(args.template) as f:
base_config = yaml.safe_load(
string.Template(f.read()).substitute(substitutions))
scenario_filter = scenario_config_exporter.scenario_filter(
scenario_name_regex=args.regex,
category=args.category,
client_language=args.client_language,
server_language=args.server_language)
scenarios = scenario_config_exporter.gen_scenarios(args.language,
scenario_filter)
uniquifiers = args.uniquifiers
if args.d:
uniquifiers.append(now_string())
annotations = parse_key_value_args(args.annotations)
configs = gen_loadtest_configs(base_config,
scenarios,
loadtest_name_prefix=args.prefix,
uniquifiers=uniquifiers,
annotations=annotations,
runs_per_test=args.runs_per_test)
configure_yaml()
with open(args.output, 'w') if args.output else sys.stdout as f:
yaml.dump_all(configs, stream=f)
def gen_loadtest_configs(
base_config: Mapping[str, Any],
base_config_clients: Iterable[Mapping[str, Any]],
base_config_servers: Iterable[Mapping[str, Any]],
scenario_name_regex: str,
language_config: scenario_config_exporter.LanguageConfig,
loadtest_name_prefix: str,
uniquifier_elements: Iterable[str],
annotations: Mapping[str, str],
instances_per_client: int = 1,
runs_per_test: int = 1) -> Iterable[Dict[str, Any]]:
"""Generates LoadTest configurations for a given language config.
The LoadTest configurations are generated as YAML objects.
"""
validate_annotations(annotations)
prefix = loadtest_name_prefix or default_prefix()
cl = safe_name(language_config.client_language or language_config.language)
sl = safe_name(language_config.server_language or language_config.language)
scenario_filter = scenario_config_exporter.scenario_filter(
scenario_name_regex=scenario_name_regex,
category=language_config.category,
client_language=language_config.client_language,
server_language=language_config.server_language)
scenarios = scenario_config_exporter.gen_scenarios(
language_config.language, scenario_filter)
for scenario in scenarios:
for run_index in gen_run_indices(runs_per_test):
uniq = (uniquifier_elements +
[run_index] if run_index else uniquifier_elements)
name = loadtest_name(prefix, scenario['name'], uniq)
scenario_str = json.dumps({'scenarios': scenario},
indent=' ') + '\n'
config = copy.deepcopy(base_config)
metadata = config['metadata']
metadata['name'] = name
if 'labels' not in metadata:
metadata['labels'] = dict()
metadata['labels']['language'] = safe_name(
language_config.language)
metadata['labels']['prefix'] = prefix
if 'annotations' not in metadata:
metadata['annotations'] = dict()
metadata['annotations'].update(annotations)
metadata['annotations'].update({
'scenario': scenario['name'],
'uniquifier': '-'.join(uniq),
})
spec = config['spec']
# Select clients with the required language.
clients = [
client for client in base_config_clients
if client['language'] == cl
]
if not clients:
raise IndexError('Client language not found in template: %s' %
cl)
# Validate config for additional client instances.
if instances_per_client > 1:
c = collections.Counter(
(client.get('name', '') for client in clients))
if max(c.values()) > 1:
raise ValueError(
('Multiple instances of multiple clients requires '
'unique names, name counts for language %s: %s') %
(cl, c.most_common()))
# Name client instances with an index starting from zero.
client_instances = []
for i in range(instances_per_client):
client_instances.extend(copy.deepcopy(clients))
for client in client_instances[-len(clients):]:
client['name'] = component_name((client.get('name',
''), str(i)))
# Set clients to named instances.
spec['clients'] = client_instances
# Select servers with the required language.
servers = copy.deepcopy([
server for server in base_config_servers
if server['language'] == sl
])
if not servers:
raise IndexError('Server language not found in template: %s' %
sl)
# Name servers with an index for consistency with clients.
for i, server in enumerate(servers):
server['name'] = component_name((server.get('name',
''), str(i)))
# Set servers to named instances.
spec['servers'] = servers
# Add driver, if needed.
if 'driver' not in spec:
spec['driver'] = dict()
# Ensure driver has language and run fields.
driver = spec['driver']
if 'language' not in driver:
driver['language'] = safe_name('c++')
if 'run' not in driver:
driver['run'] = dict()
# Name the driver with an index for consistency with workers.
# There is only one driver, so the index is zero.
if 'name' not in driver or not driver['name']:
driver['name'] = '0'
spec['scenariosJSON'] = scenario_str
yield config
def gen_loadtest_configs(
base_config: Mapping[str, Any],
base_config_clients: Iterable[Mapping[str, Any]],
base_config_servers: Iterable[Mapping[str, Any]],
scenario_name_regex: str,
language_config: scenario_config_exporter.LanguageConfig,
loadtest_name_prefix: str,
uniquifier_elements: Iterable[str],
annotations: Mapping[str, str],
runs_per_test: int = 1) -> Iterable[Dict[str, Any]]:
"""Generates LoadTest configurations for a given language config.
The LoadTest configurations are generated as YAML objects.
"""
validate_annotations(annotations)
prefix = loadtest_name_prefix or default_prefix()
cl = image_language(language_config.client_language
or language_config.language)
sl = image_language(language_config.server_language
or language_config.language)
scenario_filter = scenario_config_exporter.scenario_filter(
scenario_name_regex=scenario_name_regex,
category=language_config.category,
client_language=language_config.client_language,
server_language=language_config.server_language)
scenarios = scenario_config_exporter.gen_scenarios(
language_config.language, scenario_filter)
for scenario in scenarios:
for run_index in gen_run_indices(runs_per_test):
uniq = (uniquifier_elements +
[run_index] if run_index else uniquifier_elements)
name = loadtest_name(prefix, scenario['name'], uniq)
scenario_str = json.dumps({'scenarios': scenario}, indent=' ')
config = copy.deepcopy(base_config)
metadata = config['metadata']
metadata['name'] = name
if 'labels' not in metadata:
metadata['labels'] = dict()
metadata['labels']['language'] = language_config.language
metadata['labels']['prefix'] = prefix
if 'annotations' not in metadata:
metadata['annotations'] = dict()
metadata['annotations'].update(annotations)
metadata['annotations'].update({
'scenario': scenario['name'],
'uniquifier': '-'.join(uniq),
})
spec = config['spec']
# Select clients with the required language.
spec['clients'] = [
client for client in base_config_clients
if client['language'] == cl
]
if not spec['clients']:
raise IndexError('Client language not found in template: %s' %
cl)
# Select servers with the required language.
spec['servers'] = [
server for server in base_config_servers
if server['language'] == sl
]
if not spec['servers']:
raise IndexError('Server language not found in template: %s' %
sl)
spec['scenariosJSON'] = scenario_str
yield config