JUnit annotations

from operator import itemgetter
from functools import wraps
import unittest


BEFORE, AFTER, TEST = range(3)


# NOTE: sequence is negative for LIFO sorting
def before(f, __state=[0]):
    # Annotate function, since it is still a plain function
    f.annotate_call_when = BEFORE, __state[0]
    __state[0] -= 1

    return f

def after(f, __state=[0]):
    # Annotate function, since it is still a plain function
    f.annotate_call_when = AFTER, __state[0]
    __state[0] -= 1

    return f

def test(f):
    # Annotate function, since it is still a plain function
    f.annotate_call_when = TEST, 0

    return f


class JUnitAnnotations(type):
    def __new__(cls, name, bases, ns):
        def collect(ns):
            for name, value in ns.iteritems():
                annotation, n = getattr(
                    value, "annotate_call_when", (None, 0))
                if annotation == BEFORE:
                    call_before.append((value, n))
                elif annotation == AFTER:
                    call_after.append((value, n))
                elif annotation == TEST:
                    call_test.append((name, value))

        def test_wrapper(test, call_before, call_after):
            @wraps(test)
            def wrapped(self):
                for fixture, n in call_before:
                    fixture(self)

                test(self)

                for fixture, n in call_after:
                    fixture(self)

            return wrapped

        call_before = []
        call_after = []
        call_test = []

        # Collect all annotations
        collect(ns)
        for parent in bases:
            collect(parent.__dict__)

        # Before and After functions must be called LIFO
        # The order is correct for base classes, since base classes
        # methods calls the after and before decorators first.
        call_before.sort(key=itemgetter(1))
        call_after.sort(key=itemgetter(1))

        # Wrap each annotated Test so that Before and After fixtures are
        # executed correctly
        for name, fun in call_test:
            ns[name] = test_wrapper(
                fun, call_before, call_after)

        # Let it be
        return super(JUnitAnnotations, cls).__new__(cls, name, bases, ns)


class AnnotatedTestCase(unittest.TestCase):
    __metaclass__ = JUnitAnnotations


class TestCaseBase(AnnotatedTestCase):
    @before
    def fixture_base_before_1(self):
        print "fixture base before 1"

    @before
    def fixture_base_before_2(self):
        print "fixture base before 2"

    @after
    def fixture_base_after_1(self):
        print "fixture base after 1"

    @after
    def fixture_base_after_2(self):
        print "fixture base after 2"


class TestCase(TestCaseBase):
    @before
    def fixture_before_1(self):
        print "fixture before 1"

    @before
    def fixture_before_2(self):
        print "fixture before 2"

    @after
    def fixture_after_1(self):
        print "fixture after 1"

    @after
    def fixture_after_2(self):
        print "fixture after 2"

    @test
    def test(self):
        print "test"

    # internal
    def runTest(self):
        pass


if __name__ == '__main__':
    t = TestCase()

    t.test()