Nose plugin to time tests

timetests.py

"""This plugin provides test timings to identify which tests might be
taking the most. From this information, it might be useful to couple
individual tests nose's `--with-profile` option to profile problematic
tests.

This plugin is heavily influenced by nose's `xunit` plugin.

Add this command to the way you execute nose::

    --with-test-timer

After all tests are executed, they will be sorted in ascending order.

(c) 2011 - Mahmoud Abdelkader (http://github.com/mahmoudimus)

LICENSE:
            DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
                    Version 2, December 2004

 Copyright (C) 2004 Sam Hocevar <sam@hocevar.net>

 Everyone is permitted to copy and distribute verbatim or modified
 copies of this license document, and changing it is allowed as long
 as the name is changed.

            DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
   TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION

  0. You just DO WHAT THE FUCK YOU WANT TO.

"""

import operator
from time import time

import nose
from nose.plugins.base import Plugin
import nose.util
import csv


class TestTimer(Plugin):
    """This plugin provides test timings

    """

    name = 'test-timer'
    score = 1

    def _timeTaken(self):
        if hasattr(self, '_timer'):
            taken = time() - self._timer
        else:
            # test died before it ran (probably error in setup())
            # or success/failure added before test started probably
            # due to custom TestResult munging
            taken = 0.0
        return taken

    def options(self, parser, env):
        """Sets additional command line options."""
        super(TestTimer, self).options(parser, env)
        parser.add_option('--test-timer-output-csv',
                          type='string',
                          dest='output_csv',
                          default=None,
                          help='Path to CSV output')
        parser.add_option('--test-timer-threshold',
                          type='float',
                          dest='threshold',
                          default=None,
                          help='Only capture tests that take longer than the specified number of seconds FLOAT')

    def configure(self, options, config):
        """Configures the test timer plugin."""
        super(TestTimer, self).configure(options, config)
        self.config = config
        self.options = options
        self._timed_tests = {}

    def startTest(self, test):
        """Initializes a timer before starting a test."""
        self._timer = time()


    def report(self, stream):
        """Report the test times"""
        if not self.enabled:
            return

        writer = None
        if self.options.output_csv:
            writer = csv.writer(open(self.options.output_csv, 'w'))

        d = sorted(self._timed_tests.iteritems(), key=operator.itemgetter(1))
        reports = []
        for test, time_taken in d:
            if time_taken < self.options.threshold:
                continue
            file, module, rest = test
            test_name = 'unknown test'
            if module:
                test_name = module
                if rest:
                    test_name = '%s:%s' % (module, rest)
            if writer:
                reports.append(['%0.4f' % time_taken, test_name])
            else:
                stream.writeln("%0.4f - %s" % (time_taken, test_name))
        if writer:
            writer.writerows(reports)

    def _register_time(self, test):
        self._timed_tests[test.address()] = self._timeTaken()

    def addError(self, test, err, capt=None):
        self._register_time(test)

    def addFailure(self, test, err, capt=None, tb_info=None):
        self._register_time(test)

    def addSuccess(self, test, capt=None):
        self._register_time(test)


if __name__ == '__main__':
    nose.main(addplugins=[TestTimer()])