litsupport/modules/microbenchmark.py

'''Test module to collect google benchmark results.'''
from litsupport import shellcommand
from litsupport import testplan
import json
import lit.Test


def _mutateCommandLine(context, commandline):
    cmd = shellcommand.parse(commandline)
    cmd.arguments.append("--benchmark_format=json")
    # We need stdout outself to get the benchmark csv data.
    if cmd.stdout is not None:
        raise Exception("Rerouting stdout not allowed for microbenchmarks")
    benchfile = context.tmpBase + '.bench.json'
    cmd.stdout = benchfile
    context.microbenchfiles.append(benchfile)

    return cmd.toCommandline()


def _mutateScript(context, script):
    return testplan.mutateScript(context, script, _mutateCommandLine)


def _collectMicrobenchmarkTime(context, microbenchfiles):
    for f in microbenchfiles:
        content = context.read_result_file(context, f)
        data = json.loads(content)

        # Create a micro_result for each benchmark
        for benchmark in data['benchmarks']:
            # Name for MicroBenchmark
            name = benchmark['name']

            # Create Result object with PASS
            microBenchmark = lit.Test.Result(lit.Test.PASS)

            # Add the exec_time metric for this result
            exec_time_metric = lit.Test.toMetricValue(benchmark['cpu_time'])
            microBenchmark.addMetric('exec_time', exec_time_metric)

            # Add Micro Result
            context.micro_results[name] = microBenchmark

    # returning the number of microbenchmarks collected as a metric for the
    # base test
    return ({
        'MicroBenchmarks': lit.Test.toMetricValue(len(context.micro_results))
    })


def mutatePlan(context, plan):
    context.microbenchfiles = []
    plan.runscript = _mutateScript(context, plan.runscript)
    plan.metric_collectors.append(
        lambda context: _collectMicrobenchmarkTime(context,
                                                   context.microbenchfiles)
    )