Skip to content

Commit 0f221d0

Browse files
lisroachlisroach
authored
bpo-24412: Adds cleanUps for setUpClass and setUpModule. (pythonGH-9190)
1 parent 49fa4a9 commit 0f221d0

File tree

7 files changed

+783
-18
lines changed

7 files changed

+783
-18
lines changed

‎Doc/library/unittest.rst

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1448,6 +1448,39 @@ Test cases
14481448

14491449
.. versionadded:: 3.1
14501450

1451+
.. classmethod:: addClassCleanup(function, *args, **kwargs)
1452+
1453+
Add a function to be called after :meth:`tearDownClass` to cleanup
1454+
resources used during the test class. Functions will be called in reverse
1455+
order to the order they are added (:abbr:`LIFO (last-in, first-out)`).
1456+
They are called with any arguments and keyword arguments passed into
1457+
:meth:`addClassCleanup` when they are added.
1458+
1459+
If :meth:`setUpClass` fails, meaning that :meth:`tearDownClass` is not
1460+
called, then any cleanup functions added will still be called.
1461+
1462+
.. versionadded:: 3.8
1463+
1464+
1465+
.. classmethod:: doClassCleanups()
1466+
1467+
This method is called unconditionally after :meth:`tearDownClass`, or
1468+
after :meth:`setUpClass` if :meth:`setUpClass` raises an exception.
1469+
1470+
It is responsible for calling all the cleanup functions added by
1471+
:meth:`addCleanupClass`. If you need cleanup functions to be called
1472+
*prior* to :meth:`tearDownClass` then you can call
1473+
:meth:`doCleanupsClass` yourself.
1474+
1475+
:meth:`doCleanupsClass` pops methods off the stack of cleanup
1476+
functions one at a time, so it can be called at any time.
1477+
1478+
.. versionadded:: 3.8
1479+
1480+
1481+
1482+
1483+
14511484

14521485
.. class:: FunctionTestCase(testFunc, setUp=None, tearDown=None, description=None)
14531486

@@ -2268,6 +2301,38 @@ module will be run and the ``tearDownModule`` will not be run. If the exception
22682301
:exc:`SkipTest` exception then the module will be reported as having been skipped
22692302
instead of as an error.
22702303

2304+
To add cleanup code that must be run even in the case of an exception, use
2305+
``addModuleCleanup``:
2306+
2307+
2308+
.. function:: addModuleCleanup(function, *args, **kwargs)
2309+
2310+
Add a function to be called after :func:`tearDownModule` to cleanup
2311+
resources used during the test class. Functions will be called in reverse
2312+
order to the order they are added (:abbr:`LIFO (last-in, first-out)`).
2313+
They are called with any arguments and keyword arguments passed into
2314+
:meth:`addModuleCleanup` when they are added.
2315+
2316+
If :meth:`setUpModule` fails, meaning that :func:`tearDownModule` is not
2317+
called, then any cleanup functions added will still be called.
2318+
2319+
.. versionadded:: 3.8
2320+
2321+
2322+
.. function:: doModuleCleanups()
2323+
2324+
This function is called unconditionally after :func:`tearDownModule`, or
2325+
after :func:`setUpModule` if :func:`setUpModule` raises an exception.
2326+
2327+
It is responsible for calling all the cleanup functions added by
2328+
:func:`addCleanupModule`. If you need cleanup functions to be called
2329+
*prior* to :func:`tearDownModule` then you can call
2330+
:func:`doModuleCleanups` yourself.
2331+
2332+
:func:`doModuleCleanups` pops methods off the stack of cleanup
2333+
functions one at a time, so it can be called at any time.
2334+
2335+
.. versionadded:: 3.8
22712336

22722337
Signal Handling
22732338
---------------

‎Doc/whatsnew/3.8.rst

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,15 @@ unicodedata
233233
is in a specific normal form. (Contributed by Max Belanger and David Euresti in
234234
:issue:`32285`).
235235

236+
unittest
237+
--------
238+
239+
* Added :func:`~unittest.addModuleCleanup()` and
240+
:meth:`~unittest.TestCase.addClassCleanup()` to unittest to support
241+
cleanups for :func:`~unittest.setUpModule()` and
242+
:meth:`~unittest.TestCase.setUpClass()`.
243+
(Contributed by Lisa Roach in :issue:`24412`.)
244+
236245
venv
237246
----
238247

‎Lib/unittest/__init__.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,17 @@ def testMultiply(self):
4848
'TextTestRunner', 'TestLoader', 'FunctionTestCase', 'main',
4949
'defaultTestLoader', 'SkipTest', 'skip', 'skipIf', 'skipUnless',
5050
'expectedFailure', 'TextTestResult', 'installHandler',
51-
'registerResult', 'removeResult', 'removeHandler']
51+
'registerResult', 'removeResult', 'removeHandler',
52+
'addModuleCleanup']
5253

5354
# Expose obsolete functions for backwards compatibility
5455
__all__.extend(['getTestCaseNames', 'makeSuite', 'findTestCases'])
5556

5657
__unittest = True
5758

5859
from .result import TestResult
59-
from .case import (TestCase, FunctionTestCase, SkipTest, skip, skipIf,
60-
skipUnless, expectedFailure)
60+
from .case import (addModuleCleanup, TestCase, FunctionTestCase, SkipTest, skip,
61+
skipIf, skipUnless, expectedFailure)
6162
from .suite import BaseTestSuite, TestSuite
6263
from .loader import (TestLoader, defaultTestLoader, makeSuite, getTestCaseNames,
6364
findTestCases)

‎Lib/unittest/case.py

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,30 @@ def testPartExecutor(self, test_case, isTest=False):
8484
def _id(obj):
8585
return obj
8686

87+
88+
_module_cleanups = []
89+
def addModuleCleanup(function, *args, **kwargs):
90+
"""Same as addCleanup, except the cleanup items are called even if
91+
setUpModule fails (unlike tearDownModule)."""
92+
_module_cleanups.append((function, args, kwargs))
93+
94+
95+
def doModuleCleanups():
96+
"""Execute all module cleanup functions. Normally called for you after
97+
tearDownModule."""
98+
exceptions = []
99+
while _module_cleanups:
100+
function, args, kwargs = _module_cleanups.pop()
101+
try:
102+
function(*args, **kwargs)
103+
except Exception as exc:
104+
exceptions.append(exc)
105+
if exceptions:
106+
# Swallows all but first exception. If a multi-exception handler
107+
# gets written we should use that here instead.
108+
raise exceptions[0]
109+
110+
87111
def skip(reason):
88112
"""
89113
Unconditionally skip a test.
@@ -390,6 +414,8 @@ class TestCase(object):
390414

391415
_classSetupFailed = False
392416

417+
_class_cleanups = []
418+
393419
def __init__(self, methodName='runTest'):
394420
"""Create an instance of the class that will use the named test
395421
method when executed. Raises a ValueError if the instance does
@@ -445,6 +471,12 @@ def addCleanup(self, function, *args, **kwargs):
445471
Cleanup items are called even if setUp fails (unlike tearDown)."""
446472
self._cleanups.append((function, args, kwargs))
447473

474+
@classmethod
475+
def addClassCleanup(cls, function, *args, **kwargs):
476+
"""Same as addCleanup, except the cleanup items are called even if
477+
setUpClass fails (unlike tearDownClass)."""
478+
cls._class_cleanups.append((function, args, kwargs))
479+
448480
def setUp(self):
449481
"Hook method for setting up the test fixture before exercising it."
450482
pass
@@ -651,9 +683,21 @@ def doCleanups(self):
651683
function(*args, **kwargs)
652684

653685
# return this for backwards compatibility
654-
# even though we no longer us it internally
686+
# even though we no longer use it internally
655687
return outcome.success
656688

689+
@classmethod
690+
def doClassCleanups(cls):
691+
"""Execute all class cleanup functions. Normally called for you after
692+
tearDownClass."""
693+
cls.tearDown_exceptions = []
694+
while cls._class_cleanups:
695+
function, args, kwargs = cls._class_cleanups.pop()
696+
try:
697+
function(*args, **kwargs)
698+
except Exception as exc:
699+
cls.tearDown_exceptions.append(sys.exc_info())
700+
657701
def __call__(self, *args, **kwds):
658702
return self.run(*args, **kwds)
659703

‎Lib/unittest/suite.py

Lines changed: 50 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -166,10 +166,18 @@ def _handleClassSetUp(self, test, result):
166166
raise
167167
currentClass._classSetupFailed = True
168168
className = util.strclass(currentClass)
169-
errorName = 'setUpClass (%s)' % className
170-
self._addClassOrModuleLevelException(result, e, errorName)
169+
self._createClassOrModuleLevelException(result, e,
170+
'setUpClass',
171+
className)
171172
finally:
172173
_call_if_exists(result, '_restoreStdout')
174+
if currentClass._classSetupFailed is True:
175+
currentClass.doClassCleanups()
176+
if len(currentClass.tearDown_exceptions) > 0:
177+
for exc in currentClass.tearDown_exceptions:
178+
self._createClassOrModuleLevelException(
179+
result, exc[1], 'setUpClass', className,
180+
info=exc)
173181

174182
def _get_previous_module(self, result):
175183
previousModule = None
@@ -199,21 +207,37 @@ def _handleModuleFixture(self, test, result):
199207
try:
200208
setUpModule()
201209
except Exception as e:
210+
try:
211+
case.doModuleCleanups()
212+
except Exception as exc:
213+
self._createClassOrModuleLevelException(result, exc,
214+
'setUpModule',
215+
currentModule)
202216
if isinstance(result, _DebugResult):
203217
raise
204218
result._moduleSetUpFailed = True
205-
errorName = 'setUpModule (%s)' % currentModule
206-
self._addClassOrModuleLevelException(result, e, errorName)
219+
self._createClassOrModuleLevelException(result, e,
220+
'setUpModule',
221+
currentModule)
207222
finally:
208223
_call_if_exists(result, '_restoreStdout')
209224

210-
def _addClassOrModuleLevelException(self, result, exception, errorName):
225+
def _createClassOrModuleLevelException(self, result, exc, method_name,
226+
parent, info=None):
227+
errorName = f'{method_name} ({parent})'
228+
self._addClassOrModuleLevelException(result, exc, errorName, info)
229+
230+
def _addClassOrModuleLevelException(self, result, exception, errorName,
231+
info=None):
211232
error = _ErrorHolder(errorName)
212233
addSkip = getattr(result, 'addSkip', None)
213234
if addSkip is not None and isinstance(exception, case.SkipTest):
214235
addSkip(error, str(exception))
215236
else:
216-
result.addError(error, sys.exc_info())
237+
if not info:
238+
result.addError(error, sys.exc_info())
239+
else:
240+
result.addError(error, info)
217241

218242
def _handleModuleTearDown(self, result):
219243
previousModule = self._get_previous_module(result)
@@ -235,10 +259,17 @@ def _handleModuleTearDown(self, result):
235259
except Exception as e:
236260
if isinstance(result, _DebugResult):
237261
raise
238-
errorName = 'tearDownModule (%s)' % previousModule
239-
self._addClassOrModuleLevelException(result, e, errorName)
262+
self._createClassOrModuleLevelException(result, e,
263+
'tearDownModule',
264+
previousModule)
240265
finally:
241266
_call_if_exists(result, '_restoreStdout')
267+
try:
268+
case.doModuleCleanups()
269+
except Exception as e:
270+
self._createClassOrModuleLevelException(result, e,
271+
'tearDownModule',
272+
previousModule)
242273

243274
def _tearDownPreviousClass(self, test, result):
244275
previousClass = getattr(result, '_previousTestClass', None)
@@ -261,10 +292,19 @@ def _tearDownPreviousClass(self, test, result):
261292
if isinstance(result, _DebugResult):
262293
raise
263294
className = util.strclass(previousClass)
264-
errorName = 'tearDownClass (%s)' % className
265-
self._addClassOrModuleLevelException(result, e, errorName)
295+
self._createClassOrModuleLevelException(result, e,
296+
'tearDownClass',
297+
className)
266298
finally:
267299
_call_if_exists(result, '_restoreStdout')
300+
previousClass.doClassCleanups()
301+
if len(previousClass.tearDown_exceptions) > 0:
302+
for exc in previousClass.tearDown_exceptions:
303+
className = util.strclass(previousClass)
304+
self._createClassOrModuleLevelException(result, exc[1],
305+
'tearDownClass',
306+
className,
307+
info=exc)
268308

269309

270310
class _ErrorHolder(object):

0 commit comments

Comments
 (0)