From cca2cf3990c592bb7981fd79966449ab83b77670 Mon Sep 17 00:00:00 2001
From: herr kaste <herr.kaste@gmail.com>
Date: Fri, 7 Dec 2018 00:11:40 +0100
Subject: [PATCH 1/3] Make `doCleanups` work for the DeferrableTestCase

---
 tests/test_ensure_do_cleanups.py | 15 +++++++++++++++
 unittesting/core/st3/case.py     | 17 ++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_ensure_do_cleanups.py

diff --git a/tests/test_ensure_do_cleanups.py b/tests/test_ensure_do_cleanups.py
new file mode 100644
index 00000000..3d60c62f
--- /dev/null
+++ b/tests/test_ensure_do_cleanups.py
@@ -0,0 +1,15 @@
+from unittesting import DeferrableTestCase
+
+
+class TestDoCleanups(DeferrableTestCase):
+
+    def test_ensure_do_cleanups_works(self):
+        messages = []
+
+        def work(message):
+            messages.append(message)
+
+        self.addCleanup(work, 1)
+        yield from self.doCleanups()
+
+        self.assertEqual(messages, [1])
diff --git a/unittesting/core/st3/case.py b/unittesting/core/st3/case.py
index 9d52e8b5..c3164cc1 100644
--- a/unittesting/core/st3/case.py
+++ b/unittesting/core/st3/case.py
@@ -74,7 +74,7 @@ def run(self, result=None):
                 if isiterable(deferred):
                     yield from deferred
 
-            self.doCleanups()
+            yield from self.doCleanups()
             if outcome.success:
                 result.addSuccess(self)
             else:
@@ -110,3 +110,18 @@ def run(self, result=None):
                 stopTestRun = getattr(result, 'stopTestRun', None)
                 if stopTestRun is not None:
                     stopTestRun()
+
+    def doCleanups(self):
+        """Execute all cleanup functions.
+
+        Normally called for you after tearDown.
+        """
+        outcome = self._outcomeForDoCleanups or _Outcome()
+        while self._cleanups:
+            function, args, kwargs = self._cleanups.pop()
+            part = lambda: function(*args, **kwargs)  # noqa: E731
+            yield from self._executeTestPart(part, outcome)
+
+        # return this for backwards compatibility
+        # even though we no longer us it internally
+        return outcome.success

From 18b20a4b76d66d2faeecac04510d2f56beaa6ca7 Mon Sep 17 00:00:00 2001
From: herr kaste <herr.kaste@gmail.com>
Date: Fri, 7 Dec 2018 00:21:39 +0100
Subject: [PATCH 2/3] `_executeTestPart` is a generator and always iterable

---
 unittesting/core/st3/case.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/unittesting/core/st3/case.py b/unittesting/core/st3/case.py
index c3164cc1..29e4123c 100644
--- a/unittesting/core/st3/case.py
+++ b/unittesting/core/st3/case.py
@@ -63,16 +63,10 @@ def run(self, result=None):
             outcome = _Outcome()
             self._outcomeForDoCleanups = outcome
 
-            deferred = self._executeTestPart(self.setUp, outcome)
-            if isiterable(deferred):
-                yield from deferred
+            yield from self._executeTestPart(self.setUp, outcome)
             if outcome.success:
-                deferred = self._executeTestPart(testMethod, outcome, isTest=True)
-                if isiterable(deferred):
-                    yield from deferred
-                deferred = self._executeTestPart(self.tearDown, outcome)
-                if isiterable(deferred):
-                    yield from deferred
+                yield from self._executeTestPart(testMethod, outcome, isTest=True)
+                yield from self._executeTestPart(self.tearDown, outcome)
 
             yield from self.doCleanups()
             if outcome.success:

From cbbf178a59561e828214ff88e0c73ec0716fa926 Mon Sep 17 00:00:00 2001
From: herr kaste <herr.kaste@gmail.com>
Date: Fri, 7 Dec 2018 23:47:40 +0100
Subject: [PATCH 3/3] Test implicit `doCleanups` on tearDown

---
 tests/test_ensure_do_cleanups.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/tests/test_ensure_do_cleanups.py b/tests/test_ensure_do_cleanups.py
index 3d60c62f..6a827bd0 100644
--- a/tests/test_ensure_do_cleanups.py
+++ b/tests/test_ensure_do_cleanups.py
@@ -1,9 +1,9 @@
 from unittesting import DeferrableTestCase
 
 
-class TestDoCleanups(DeferrableTestCase):
+class TestExplicitDoCleanups(DeferrableTestCase):
 
-    def test_ensure_do_cleanups_works(self):
+    def test_manually_calling_do_cleanups_works(self):
         messages = []
 
         def work(message):
@@ -13,3 +13,14 @@ def work(message):
         yield from self.doCleanups()
 
         self.assertEqual(messages, [1])
+
+
+cleanup_called = []
+
+
+class TestImplicitDoCleanupsOnTeardown(DeferrableTestCase):
+    def test_a_prepare(self):
+        self.addCleanup(lambda: cleanup_called.append(1))
+
+    def test_b_assert(self):
+        self.assertEqual(cleanup_called, [1])