Python Decorator that automatically unstubs all Mockito stubs after each test case

use_mocking.py

from unittest2 import TestCase
from mockito import when, unstub
from uuid import UUID
import uuid


def use_mocking(test_suite):
    wrapped_tearDown = test_suite.tearDown if 'tearDown' in dir(test_suite) else lambda self: True

    def tearDown(self):
        try:
            wrapped_tearDown(self)
        except:
            raise
        finally:
            unstub()

    test_suite.tearDown = tearDown
    return test_suite


class Person(object):
    def __init__(self, first_name, last_name):
        self.id = uuid.uuid4()
        self.first_name = first_name
        self.last_name = last_name

    def hash(self):
        return '{id}-{initials}'.format(id=self.id.hex, initials=self._initials())

    def _initials(self):
        def initial(name):
            return name[0].upper()

        return '{0}{1}'.format(initial(self.first_name), initial(self.last_name))


@use_mocking
class PersonTest(TestCase):
    #tearDown is created with a call to unstub()
    #or, if a tearDown method exists, unstub is called after calling tearDown

    def test_hash(self):
        id = UUID('00000000-0000-0000-0000-000000000000')
        when(uuid).uuid4().thenReturn(id)

        person = Person('John', 'Napier')
        self.assertEquals('00000000000000000000000000000000-JN', person.hash())
        #after execution, uuid.uuid4 is automatically unstubbed

    def test_uuid(self):
        #uuid.uuid4() is random again after being unstubbed after previous test case
        self.assertNotEquals(uuid.uuid4(), uuid.uuid4())