Experiment in tracing tests back to source they execute

testtrace.py

#!/usr/bin/env python

from __future__ import absolute_import, print_function, unicode_literals

from collections import defaultdict, namedtuple
from gettext import gettext as _, ngettext
from unittest import TestCase
from unittest.loader import TestLoader
from unittest.result import TestResult
from unittest.suite import TestSuite
import sys


class AnimalBase:

    SOUND = None

    @classmethod
    def get_sound(cls):
        return cls.SOUND


class Dog(AnimalBase):

    SOUND = "bark bark"

    def bark(self):
        return self.get_sound()


class Cat:

    SOUND = "meeeow!"

    def meow(self):
        return self.get_sound()


class DogTests(TestCase):
    """ Tests for the Dog class """

    def test_barking(self):
        """ Check if the dog can bark """
        # This is artificially verbose so that we can see multiple lines being
        # traced by the runtime trace monitor.
        dog = Dog()
        expected = "bark bark"
        actual = dog.bark()
        self.assertEqual(expected, actual)


class CatTests(TestCase):
    """ Tests for the Cat class """

    def test_meowing(self):
        """ Check if the cat can meow """
        cat = Cat()
        expected = "meeeow!"
        actual = cat.meow()
        self.assertEqual(expected, actual)


def trace(frame, event, arg):
    if event == 'call':
        # If a function is being called keep tracing execution
        # inside the function but do that only for things not
        # in the standard library.
        # Skip anything in standard library
        if not frame.f_code.co_filename.startswith(sys.prefix):
            print("TRACE: frame:{frame} event:{event} arg:{arg}".format(
                frame=frame, event=event, arg=arg))
            return trace
        else:
            return
    elif event == 'line':
        # Print the instruction offset for each line of the function
        print("TRACE: location: {file}:{line} {name}()".format(
            file=frame.f_code.co_filename,
            line=frame.f_lineno,
            name=frame.f_code.co_name))


CodeLocation = namedtuple("CodeLocation", "filename lineno name")


class TracingTestResult(TestResult):
    """
    A TestResult subclass that traces execution to know
    which lines were executed by a particular test.
    """

    def __init__(self, stream=None, descriptions=None, verbosity=None,
                 trace_details=False):
        super(TracingTestResult, self).__init__(
            stream, descriptions, verbosity)
        self._location_to_tests = defaultdict(set)
        self._test_to_locations = defaultdict(set)
        self._current_test = None
        self._trace_details = trace_details

    def startTest(self, test):
        super(TracingTestResult, self).startTest(test)
        # print("Starting test: {test}".format(test=test))
        self._current_test = test
        sys.settrace(self._trace)

    def stopTest(self, test):
        sys.settrace(None)
        super(TracingTestResult, self).stopTest(test)
        # print("Stopped test: {test}".format(test=test))
        self._current_test = None

    def _trace(self, frame, event, arg):
        # Skip stuff that happens in standard library
        if frame.f_code.co_filename.startswith(sys.prefix):
            return
        if event == 'call':
            self._associate_with_location(frame)
            # If requested, trace with per-line precision
            if self._trace_details:
                return self._trace
        elif event == 'line':
            self._associate_with_location(frame)

    def _associate_with_location(self, frame):
        location = CodeLocation(
            frame.f_code.co_filename, frame.f_lineno, frame.f_code.co_name)
        # Associate the location with the current test
        self._location_to_tests[location].add(self._current_test.id())
        self._test_to_locations[self._current_test.id()].add(location)


def main():
    loader = TestLoader()
    suite = TestSuite()
    suite.addTests(loader.loadTestsFromTestCase(DogTests))
    suite.addTests(loader.loadTestsFromTestCase(CatTests))
    result = TracingTestResult()
    suite.run(result)
    dump_suite(suite, result)


def dump_suite(suite, result):
    assert isinstance(suite, TestSuite)
    assert isinstance(result, TracingTestResult)
    print(ngettext(
        "There is {0} test",
        "There are {0} tests",
        suite.countTestCases()
    ).format(suite.countTestCases()))
    for index, test in enumerate(suite, 1):
        print(_(
            "{index}: {id!r}, {test}"
        ).format(index=index, id=test.id(), test=test))
        print("Referenced locations")
        for location in result._test_to_locations[test.id()]:
            print(" - {0}".format(location))


if __name__ == "__main__":
    main()