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, for use by annotated class decorator
    f.annotate_call_when = BEFORE, __state[0]
    __state[0] -= 1

    return f

def after(f, __state=[0]):
    # Annotate function, for use by annotated class decorator
    f.annotate_call_when = AFTER, __state[0]
    __state[0] -= 1

    return f

def test(f):
    # Annotate function, for use by annotated class decorator
    f.annotate_call_when = TEST, 0

    return f

def annotated(cls):
    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
    for base in cls.__mro__:
        collect(base.__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:
       test = test_wrapper(fun, call_before, call_after)
       setattr(cls, name, test)

    return cls



#
# Example
#

@annotated
class TestCaseBase(unittest.TestCase):
    @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"


@annotated
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()