summaryrefslogtreecommitdiffstats
path: root/Lib/unittest/case.py
diff options
context:
space:
mode:
Diffstat (limited to 'Lib/unittest/case.py')
-rw-r--r--Lib/unittest/case.py15
1 files changed, 14 insertions, 1 deletions
diff --git a/Lib/unittest/case.py b/Lib/unittest/case.py
index 758924d..811f5df 100644
--- a/Lib/unittest/case.py
+++ b/Lib/unittest/case.py
@@ -442,12 +442,25 @@ class TestCase(object):
"""
self._type_equality_funcs[typeobj] = function
- def addCleanup(self, function, *args, **kwargs):
+ def addCleanup(*args, **kwargs):
"""Add a function, with arguments, to be called when the test is
completed. Functions added are called on a LIFO basis and are
called after tearDown on test failure or success.
Cleanup items are called even if setUp fails (unlike tearDown)."""
+ if len(args) >= 2:
+ self, function, *args = args
+ elif not args:
+ raise TypeError("descriptor 'addCleanup' of 'TestCase' object "
+ "needs an argument")
+ elif 'function' in kwargs:
+ function = kwargs.pop('function')
+ self, *args = args
+ else:
+ raise TypeError('addCleanup expected at least 1 positional '
+ 'argument, got %d' % (len(args)-1))
+ args = tuple(args)
+
self._cleanups.append((function, args, kwargs))
def setUp(self):