diff -r f954ee489896 Lib/unittest/case.py --- a/Lib/unittest/case.py Sun Jul 15 16:12:54 2012 +0100 +++ b/Lib/unittest/case.py Mon Jul 16 13:40:12 2012 -0700 @@ -2,6 +2,7 @@ import sys import functools +import contextlib import difflib import pprint import re @@ -371,6 +372,10 @@ return "<%s testMethod=%s>" % \ (strclass(self.__class__), self._testMethodName) + @property + def _testMethod(self): + return getattr(self, self._testMethodName) + def _addSkip(self, result, reason): addSkip = getattr(result, 'addSkip', None) if addSkip is not None: @@ -410,6 +415,32 @@ outcome.success = False outcome.errors.append(sys.exc_info()) + @contextlib.contextmanager + def setUpContext(self): + outcome = self._outcomeForDoCleanups + self._executeTestPart(self.setUp, outcome) + setUpSuccess = outcome.success + + yield + + if setUpSuccess: + self._executeTestPart(self.tearDown, outcome) + + self.doCleanups() + + def runTestMethod(self): + outcome = self._outcomeForDoCleanups + if outcome.success: + self._executeTestPart(self._testMethod, outcome, isTest=True) + + def executeTest(self): + """Execute the test with setUp, tearDown, and Cleanups. + + This method is meant to be overridden. + """ + with self.setUpContext(): + self.runTestMethod() + def run(self, result=None): orig_result = result if result is None: @@ -420,7 +451,7 @@ result.startTest(self) - testMethod = getattr(self, self._testMethodName) + testMethod = self._testMethod if (getattr(self.__class__, "__unittest_skip__", False) or getattr(testMethod, "__unittest_skip__", False)): # If the class or method was skipped. @@ -435,12 +466,8 @@ outcome = _Outcome() self._outcomeForDoCleanups = outcome - self._executeTestPart(self.setUp, outcome) - if outcome.success: - self._executeTestPart(testMethod, outcome, isTest=True) - self._executeTestPart(self.tearDown, outcome) + self.executeTest() - self.doCleanups() if outcome.success: result.addSuccess(self) else: