#!/usr/bin/env python import toml import jinja2 import json import argparse import os import pathlib import subprocess import time import analyze import re TESTGROUND_BIN = 'testground' DEFAULT_GS_VERSION = 'latest' # setting this build tag lets us compile test code that targets the new API from the hardening branch HARDENED_API_BUILD_TAG = 'hardened_api' # Testground build/run config settings to use when running on kubernetes K8S_BUILD_CONFIG = {'bypass_cache': True, 'push_registry': True, 'registry_type': 'aws'} K8S_RUN_CONFIG = {} def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('param_files', nargs='*', help='name of one or more parameter files to use when generating composition from template. ' + 'if a param is defined in multiple files, last one wins.') parser.add_argument('--name', help='name of composition. will be used to create output directory.') parser.add_argument('--template_dir', default='./templates/baseline', help='path to directory containing composition template and param files') parser.add_argument('-o', '--output', help='directory to write composition file and test outputs to', default='./output') parser.add_argument('--branch', default='master', help='configures the test to use the API from the given gossipsub branch') parser.add_argument('--commit', help='configures the test to use the API from the given gossipsub commit') parser.add_argument('--k8s', action='store_true', default=False, help='runs the test on kubernetes') parser.add_argument('--dry-run', dest='dry_run', action='store_true', default=False, help='skip running tests, just write composition files and exit') parser.add_argument('--instances', type=int, help='override the total number of test instances. equivalent to -D N_NODES=x') parser.add_argument('-D', '--define', dest='definitions', action='append', metavar='', help='set template variable `key` to `value`, e.g. -D T_RUN=10m') parser.add_argument('--skip-analysis', dest='skip_analysis', action='store_true', default=False, help='skip analysis phase after test run (can run manually later with analyze.py)') return parser.parse_args() def get_param_filepath(template_dir, param_filename): p = param_filename if os.path.exists(p): return p p = os.path.join(template_dir, 'params', p) if os.path.exists(p): return p p = param_filename + '.toml' if os.path.exists(p): return p p = os.path.join(template_dir, 'params', p) if os.path.exists(p): return p raise ValueError("can't find param file " + param_filename) def load_params(template_dir, param_files): base_params_path = os.path.join(template_dir, 'params', '_base.toml') paths = [get_param_filepath(template_dir, p) for p in param_files] params = dict() for path in [base_params_path] + paths: p = toml.load(path) for k, v in p.items(): params[k] = v return params # The TOPOLOGY parameter can be either # - a JSON representation of a topology # - a path to a file in that format def parse_topology(params): if 'TOPOLOGY' not in params: params['TOPOLOGY'] = None return # If it's already JSON, we're all set top = params['TOPOLOGY'] if len(top) > 0 and top[0] == '{': return # It's not JSON so assume it's a file path if not os.path.exists(top): raise ValueError("can't find topology file " + top) # Make sure the file contents is JSON with open(top, 'r') as infile: contents = infile.read() jsonstr = json.loads(contents) params['TOPOLOGY'] = jsonstr # N_CONTAINER_NODES_TOTAL is the total number of nodes including multiple nodes # per container def parse_n_container_nodes_total(params): n_nodes = params.get('N_NODES', 20) n_nodes_cont_honest = params.get('N_HONEST_PEERS_PER_NODE', 1) total = n_nodes * n_nodes_cont_honest params['N_CONTAINER_NODES_TOTAL'] = total def render_template(template_dir, params): env = jinja2.Environment( loader=jinja2.FileSystemLoader(template_dir) ) template = env.get_template('template.toml.j2') return template.render(**params) def composition_name(): ts = time.strftime("%Y%m%d-%H%M%S") return 'pubsub-test-{}'.format(ts) def mkdirp(dirpath): pathlib.Path(dirpath).mkdir(parents=True, exist_ok=True) def run_composition(comp_filepath, output_dir, k8s=False): archive_type = 'tgz' if k8s else 'zip' outfilename = 'test-output.{}'.format(archive_type) outpath = os.path.join(output_dir, outfilename) print('running testground composition {}'.format(comp_filepath)) print("writing test outputs to {}".format(output_dir)) cmd = [TESTGROUND_BIN, 'run', 'composition', '-f', comp_filepath, '--collect', '-o', outpath] subprocess.run(cmd, check=True) print('test completed successfully!') return outpath def pubsub_commit(ref_str): # if the input looks like a git commit already, just return it as-is if re.match(r'\b([a-f0-9]{40})\b', ref_str): return ref_str out = subprocess.run(['git', 'ls-remote', 'git://github.com/libp2p/go-libp2p-pubsub'], check=True, capture_output=True, text=True) # look for matching branch or tag in output pattern = r'^\b([a-f0-9]{40})\b.*refs/(heads|tags)/' + ref_str + '$' m = re.search(pattern, out.stdout, re.MULTILINE) if not m: raise ValueError('no branch or tag found matching {}'.format(ref_str)) return m.group(1) def run(): args = parse_args() template_dir = args.template_dir params = load_params(template_dir, args.param_files) branch = None if args.branch: branch = args.branch params['GS_VERSION'] = pubsub_commit(args.branch) if args.commit: params['GS_VERSION'] = args.commit if branch is None and args.commit is None: params['GS_VERSION'] = 'latest' gs_version_msg = 'Using go-libp2p-pubsub commit ' + params['GS_VERSION'] if branch: gs_version_msg += ' (' + branch + ')' print(gs_version_msg) if args.k8s: params['TEST_RUNNER'] = 'cluster:k8s' params['BUILD_CONFIG'] = toml.dumps(K8S_BUILD_CONFIG) params['RUN_CONFIG'] = toml.dumps(K8S_RUN_CONFIG) if args.instances is not None: params['N_NODES'] = args.instances if args.definitions is not None: for defstr in args.definitions: [k, v] = defstr.split('=') try: v = float(v) if v.is_integer(): v = int(v) params[k] = v except: params[k] = v parse_topology(params) parse_n_container_nodes_total(params) comp = composition_name() if args.name: comp += '-' + args.name workdir = os.path.join(args.output, comp) pathlib.Path(workdir).mkdir(parents=True, exist_ok=True) comp_filepath = os.path.join(workdir, 'composition.toml') with open(comp_filepath, 'w', encoding='utf8') as f: f.write(render_template(template_dir, params)) param_filepath = os.path.join(workdir, 'template-params.toml') with open(param_filepath, 'wt') as f: toml.dump(params, f) print('wrote composition file to {}'.format(comp_filepath)) if args.dry_run: print('dry run. skipping test execution') return test_output_archive = run_composition(comp_filepath, workdir, args.k8s) if not args.skip_analysis: print('extracting test output data') analyze.extract_test_outputs(test_output_archive) if __name__ == "__main__": run()