1 files changed, 56 insertions, 11 deletions
diff --git a/tests/util/test_logcontext.py b/tests/util/test_logcontext.py
index 4850722bc5..7707fbb1f0 100644
--- a/tests/util/test_logcontext.py
+++ b/tests/util/test_logcontext.py
@@ -36,24 +36,28 @@ class LoggingContextTestCase(unittest.TestCase):
             yield sleep(0)
             self._check_test_key("one")
 
-    def _test_preserve_fn(self, function):
+    def _test_run_in_background(self, function):
         sentinel_context = LoggingContext.current_context()
 
         callback_completed = [False]
 
-        @defer.inlineCallbacks
-        def cb():
+        def test():
             context_one.request = "one"
-            yield function()
-            self._check_test_key("one")
+            d = function()
 
-            callback_completed[0] = True
+            def cb(res):
+                self._check_test_key("one")
+                callback_completed[0] = True
+                return res
+            d.addCallback(cb)
+
+            return d
 
         with LoggingContext() as context_one:
             context_one.request = "one"
 
             # fire off function, but don't wait on it.
-            logcontext.preserve_fn(cb)()
+            logcontext.run_in_background(test)
 
             self._check_test_key("one")
 
@@ -80,20 +84,31 @@ class LoggingContextTestCase(unittest.TestCase):
         # test is done once d2 finishes
         return d2
 
-    def test_preserve_fn_with_blocking_fn(self):
+    def test_run_in_background_with_blocking_fn(self):
         @defer.inlineCallbacks
         def blocking_function():
             yield sleep(0)
 
-        return self._test_preserve_fn(blocking_function)
+        return self._test_run_in_background(blocking_function)
 
-    def test_preserve_fn_with_non_blocking_fn(self):
+    def test_run_in_background_with_non_blocking_fn(self):
         @defer.inlineCallbacks
         def nonblocking_function():
             with logcontext.PreserveLoggingContext():
                 yield defer.succeed(None)
 
-        return self._test_preserve_fn(nonblocking_function)
+        return self._test_run_in_background(nonblocking_function)
+
+    @unittest.DEBUG
+    def test_run_in_background_with_chained_deferred(self):
+        # a function which returns a deferred which looks like it has been
+        # called, but is actually paused
+        def testfunc():
+            return logcontext.make_deferred_yieldable(
+                _chained_deferred_function()
+            )
+
+        return self._test_run_in_background(testfunc)
 
     @defer.inlineCallbacks
     def test_make_deferred_yieldable(self):
@@ -119,6 +134,22 @@ class LoggingContextTestCase(unittest.TestCase):
             self._check_test_key("one")
 
     @defer.inlineCallbacks
+    def test_make_deferred_yieldable_with_chained_deferreds(self):
+        sentinel_context = LoggingContext.current_context()
+
+        with LoggingContext() as context_one:
+            context_one.request = "one"
+
+            d1 = logcontext.make_deferred_yieldable(_chained_deferred_function())
+            # make sure that the context was reset by make_deferred_yieldable
+            self.assertIs(LoggingContext.current_context(), sentinel_context)
+
+            yield d1
+
+            # now it should be restored
+            self._check_test_key("one")
+
+    @defer.inlineCallbacks
     def test_make_deferred_yieldable_on_non_deferred(self):
         """Check that make_deferred_yieldable does the right thing when its
         argument isn't actually a deferred"""
@@ -132,3 +163,17 @@ class LoggingContextTestCase(unittest.TestCase):
             r = yield d1
             self.assertEqual(r, "bum")
             self._check_test_key("one")
+
+
+# a function which returns a deferred which has been "called", but
+# which had a function which returned another incomplete deferred on
+# its callback list, so won't yet call any other new callbacks.
+def _chained_deferred_function():
+    d = defer.succeed(None)
+
+    def cb(res):
+        d2 = defer.Deferred()
+        reactor.callLater(0, d2.callback, res)
+        return d2
+    d.addCallback(cb)
+    return d