diff --git a/cr8/engine.py b/cr8/engine.py index b78bdb7..69e2ba8 100644 --- a/cr8/engine.py +++ b/cr8/engine.py @@ -69,7 +69,7 @@ def _generate_statements(stmt, args, iterations, duration): class Runner: - def __init__(self, hosts, concurrency, sample_mode, session_settings): + def __init__(self, hosts, concurrency, sample_mode, session_settings=None): self.concurrency = concurrency self.client = client(hosts, session_settings=session_settings, concurrency=concurrency) self.sampler = get_sampler(sample_mode) diff --git a/cr8/run_spec.py b/cr8/run_spec.py index 99d9efa..521494c 100644 --- a/cr8/run_spec.py +++ b/cr8/run_spec.py @@ -85,6 +85,7 @@ def __init__(self, session_settings): self.benchmark_hosts = benchmark_hosts self.sample_mode = sample_mode + self.session_settings = session_settings self.spec_dir = spec_dir self.client = clients.client(benchmark_hosts, session_settings) self.result_client = clients.client(result_hosts, session_settings) @@ -180,7 +181,7 @@ def _skip_message(self, min_version, stmt): server_version='.'.join((str(x) for x in self.server_version))) return msg - def run_queries(self, queries: Iterable[dict], meta=None, session_settings=None): + def run_queries(self, queries: Iterable[dict], meta=None): for query in queries: stmt = query['statement'] iterations = query.get('iterations', 1) @@ -205,7 +206,7 @@ def run_queries(self, queries: Iterable[dict], meta=None, session_settings=None) f' Concurrency: {concurrency}\n' f' {mode_desc}: {duration or iterations}') ) - with Runner(self.benchmark_hosts, concurrency, self.sample_mode, session_settings) as runner: + with Runner(self.benchmark_hosts, concurrency, self.sample_mode, self.session_settings) as runner: if warmup > 0: runner.warmup(stmt, warmup, concurrency, args) timed_stats = runner.run( @@ -269,7 +270,7 @@ def do_run_spec(spec, queries = (q for q in spec.queries if 'name' in q and rex.match(q['name'])) else: queries = spec.queries - executor.run_queries(queries, spec.meta, spec.session_settings) + executor.run_queries(queries, spec.meta) finally: if not action or 'teardown' in action: log.info('# Running tearDown')