Repository: itamarst/eliot

Issue:
Generators (and generator-based APIs like inlineCallbacks) don't work well with current implicit context
See #201 for example.

@jonathanj was talking about an explicit context API, where you just pass it around. That would be one solution.

Another is a decorator to be added to generators that deals with reinstating correct context.

At a minimum some documentation is probably called for.

Discussion:

Patch:
diff --git a/docs/source/generating/actions.rst b/docs/source/generating/actions.rst
index 45988e1..fca785c 100644
--- a/docs/source/generating/actions.rst
+++ b/docs/source/generating/actions.rst
@@ -132,6 +132,31 @@ You can add fields to both the start message and the success message of an actio
 If you want to include some extra information in case of failures beyond the exception you can always log a regular message with that information.
 Since the message will be recorded inside the context of the action its information will be clearly tied to the result of the action by the person (or code!) reading the logs later on.
 
+Using Generators
+----------------
+
+Generators (functions with ``yield``) and context managers (``with X:``) don't mix well in Python.
+So if you're going to use ``with start_action()`` in a generator, just make sure it doesn't wrap a ``yield`` and you'll be fine.
+
+Here's what you SHOULD NOT DO:
+
+.. code-block:: python
+
+   def generator():
+       with start_action(action_type="x"):
+           # BAD! DO NOT yield inside a start_action() block:
+           yield make_result()
+
+Here's what can do instead:
+
+.. code-block:: python
+
+   def generator():
+       with start_action(action_type="x"):
+           result = make_result()
+       # This is GOOD, no yield inside the start_action() block:
+       yield result
+
 
 Non-Finishing Contexts
 ----------------------
diff --git a/docs/source/generating/twisted.rst b/docs/source/generating/twisted.rst
index e15ffa8..c1c2421 100644
--- a/docs/source/generating/twisted.rst
+++ b/docs/source/generating/twisted.rst
@@ -63,6 +63,41 @@ Logging Failures
             d.addErrback(writeFailure)
 
 
+Actions and inlineCallbacks
+---------------------------
+
+Eliot provides a decorator that is compatible with Twisted's ``inlineCallbacks`` but which also behaves well with Eliot's actions.
+Simply substitute ``eliot.twisted.inline_callbacks`` for ``twisted.internet.defer.inlineCallbacks`` in your code.
+
+To understand why, consider the following example:
+
+.. code-block:: python
+
+     from eliot import start_action
+     from twisted.internet.defer import inlineCallbacks
+
+     @inlineCallbacks  # don't use this in real code, use eliot.twisted.inline_callbacks
+     def go():
+         with start_action(action_type=u"yourapp:subsystem:frob"):
+	     d = some_deferred_api()
+	     x = yield d
+	     Message.log(message_type=u"some-report", x=x)
+
+The action started by this generator remains active as ``yield d`` gives up control to the ``inlineCallbacks`` controller.
+The next bit of code to run will be considered to be a child of ``action``.
+Since that code may be any arbitrary code that happens to get scheduled,
+this is certainly wrong.
+
+Additionally,
+when the ``inlineCallbacks`` controller resumes the generator,
+it will most likely do so with no active action at all.
+This means that the log message following the yield will be recorded with no parent action,
+also certainly wrong.
+
+These problems are solved by using ``eliot.twisted.inline_callbacks`` instead of ``twisted.internet.defer.inlineCallbacks``.
+The behavior of the two decorators is identical except that Eliot's version will preserve the generator's action context and contain it within the generator.
+This extends the ``inlineCallbacks`` illusion of "synchronous" code to Eliot actions.
+
 Actions and Deferreds
 ---------------------
 
diff --git a/docs/source/news.rst b/docs/source/news.rst
index 40e7d75..e563c7b 100644
--- a/docs/source/news.rst
+++ b/docs/source/news.rst
@@ -6,17 +6,19 @@ What's New
 
 Documentation:
 
-* Eliot has an API for testing your logs were output correctly. Until now, however, the documentation was overly focused on requiring usage of types, which are optional, so it has been rewritten to be more generic: :doc:`read more about the testing API here<generating/testing>`.
+* Eliot has an API for testing that your logs were output correctly. Until now, however, the documentation was overly focused on requiring usage of types, which are optional, so it has been rewritten to be more generic: :doc:`read more about the testing API here<generating/testing>`.
 
 Features:
 
 * Generating messages is much faster.
-* ``eliot.ValidationError``, as raised by e.g. ``capture_logging``, is now part of the public API. Fixed issue #146.
-* ``eliot.twisted.DeferredContext.addCallbacks`` now supports omitting the errback, for compatibility with Twisted's ``Deferred``. Thanks to Jean-Paul Calderone for the fix. Fixed issue #366.
 * The testing infrastructure now has slightly more informative error messages. Thanks to Jean-Paul Calderone for the bug report. Fixes issue #373.
-* ``@validate_logging`` and ``@capture_logging`` now make it clearer what caused validation errors by printing the original traceback. Thanks to Jean-Paul Calderone for the bug report. Fixes issue #365.
 * Added lower-level testing infrastructure—``eliot.testing.swap_logger`` and ``eliot.testing.check_for_errors``—which is useful for cases when the ``@capture_logging`` decorator is insufficient. For example, test methods that are async, or return Twisted ``Deferred``. See the :doc:`testing documentation<generating/testing>` for details. Thanks to Jean-Paul Calderone for the feature request. Fixes #364.
-* The testing API now has slightly more informative error messages. Thanks to Jean-Paul Calderone for the bug report. Fixes issue #373.
+* ``eliot.ValidationError``, as raised by e.g. ``capture_logging``, is now part of the public API. Fixed issue #146.
+
+Twisted-related features:
+
+* New decorator, ``@eliot.twisted.inline_callbacks`` , which is like Twisted's ``inlineCallbacks`` but which also manages the Eliot context. Thanks to Jean-Paul Calderone for the fix. Fixed issue #259.
+* ``eliot.twisted.DeferredContext.addCallbacks`` now supports omitting the errback, for compatibility with Twisted's ``Deferred``. Thanks to Jean-Paul Calderone for the fix. Fixed issue #366.
 
 Bug fixes:
 
@@ -26,6 +28,7 @@ Bug fixes:
   is now thread-safe. Thanks to Jean-Paul Calderone for the patch. Fixes issue
   #382.
 
+
 1.6.0
 ^^^^^
 
diff --git a/eliot/_action.py b/eliot/_action.py
index 9cb2733..b2b81f8 100644
--- a/eliot/_action.py
+++ b/eliot/_action.py
@@ -65,7 +65,8 @@ class _ExecutionContext(threading.local):
 
     def _get_stack(self):
         """
-        Get the stack for the current context.
+        Get the stack for the current thread, or the sub-stack if sub-stacks
+        are enabled.
         """
         stack = self.get_sub_context()
         if stack is None:
diff --git a/eliot/_generators.py b/eliot/_generators.py
new file mode 100644
index 0000000..e663a2c
--- /dev/null
+++ b/eliot/_generators.py
@@ -0,0 +1,177 @@
+"""
+Support for maintaining an action context across generator suspension.
+"""
+
+from __future__ import unicode_literals, absolute_import
+
+from sys import exc_info
+from functools import wraps
+from contextlib import contextmanager
+from weakref import WeakKeyDictionary
+
+from ._action import _context_owner, _ExecutionContext
+from . import Message
+
+
+class _GeneratorContext(object):
+    """Generator sub-context for C{_ExecutionContext}."""
+
+    def __init__(self, execution_context):
+        self._execution_context = execution_context
+        self._contexts = WeakKeyDictionary()
+        self._current_generator = None
+
+    def init_stack(self, generator):
+        """Create a new stack for the given generator."""
+        stack = list(self._execution_context._get_stack())
+        self._contexts[generator] = stack
+
+    def get_stack(self):
+        """Return the sub-stack for the current generator."""
+        if self._current_generator is None:
+            # If there is no currently active generator then we have no
+            # special stack to supply.  Let the execution context figure out a
+            # different answer on its own.
+            return None
+        # Otherwise, give back the action context stack we've been tracking
+        # for the currently active generator.  It must have been previously
+        # initialized (it's too late to do it now)!
+        return self._contexts[self._current_generator]
+
+    @contextmanager
+    def in_generator(self, generator):
+        """Context manager: set the given generator as the current generator."""
+        previous_generator = self._current_generator
+        try:
+            self._current_generator = generator
+            yield
+        finally:
+            self._current_generator = previous_generator
+
+
+class GeneratorExecutionContext(_ExecutionContext):
+    """Generator-specific C{_ExecutionContext} subclass."""
+
+    def __init__(self):
+        """This will run per-thread!"""
+        _ExecutionContext.__init__(self)
+        self.generator_context = _GeneratorContext(self)
+        self.get_sub_context = self.generator_context.get_stack
+
+
+def use_generator_context():
+    """
+    Make L{eliot_friendly_generator_function} work correctly.
+    """
+    _context_owner.set(GeneratorExecutionContext)
+
+
+def _installed():
+    return isinstance(_context_owner.context, GeneratorExecutionContext)
+
+
+class GeneratorSupportNotEnabled(Exception):
+    """
+    An attempt was made to use a decorated generator without first turning on
+    the generator context manager.
+    """
+
+
+def eliot_friendly_generator_function(original):
+    """
+    Decorate a generator function so that the Eliot action context is
+    preserved across ``yield`` expressions.
+    """
+    @wraps(original)
+    def wrapper(*a, **kw):
+        # This isn't going to work if you don't have the generator context
+        # manager installed.
+        if not _installed():
+            raise GeneratorSupportNotEnabled()
+
+        # Keep track of whether the next value to deliver to the generator is
+        # a non-exception or an exception.
+        ok = True
+
+        # Keep track of the next value to deliver to the generator.
+        value_in = None
+
+        # Create the generator with a call to the generator function.  This
+        # happens with whatever Eliot action context happens to be active,
+        # which is fine and correct and also irrelevant because no code in the
+        # generator function can run until we call send or throw on it.
+        gen = original(*a, **kw)
+
+        # Initialize the per-generator Eliot action context stack to the
+        # current action stack.  This might be the main stack or, if another
+        # decorated generator is running, it might be the stack for that
+        # generator.  Not our business.
+        generator_context = _context_owner.context.generator_context
+        generator_context.init_stack(gen)
+        while True:
+            try:
+                # Whichever way we invoke the generator, we will do it
+                # with the Eliot action context stack we've saved for it.
+                # Then the context manager will re-save it and restore the
+                # "outside" stack for us.
+                #
+                # Regarding the support of Twisted's inlineCallbacks-like
+                # functionality (see eliot.twisted.inline_callbacks):
+                #
+                # The invocation may raise the inlineCallbacks internal
+                # control flow exception _DefGen_Return.  It is not wrong to
+                # just let that propagate upwards here but inlineCallbacks
+                # does think it is wrong.  The behavior triggers a
+                # DeprecationWarning to try to get us to fix our code.  We
+                # could explicitly handle and re-raise the _DefGen_Return but
+                # only at the expense of depending on a private Twisted API.
+                # For now, I'm opting to try to encourage Twisted to fix the
+                # situation (or at least not worsen it):
+                # https://twistedmatrix.com/trac/ticket/9590
+                #
+                # Alternatively, _DefGen_Return is only required on Python 2.
+                # When Python 2 support is dropped, this concern can be
+                # eliminated by always using `return value` instead of
+                # `returnValue(value)` (and adding the necessary logic to the
+                # StopIteration handler below).
+                with generator_context.in_generator(gen):
+                    if ok:
+                        value_out = gen.send(value_in)
+                    else:
+                        value_out = gen.throw(*value_in)
+                    # We have obtained a value from the generator.  In
+                    # giving it to us, it has given up control.  Note this
+                    # fact here.  Importantly, this is within the
+                    # generator's action context so that we get a good
+                    # indication of where the yield occurred.
+                    #
+                    # This is noisy, enable only for debugging:
+                    if wrapper.debug:
+                        Message.log(message_type=u"yielded")
+            except StopIteration:
+                # When the generator raises this, it is signaling
+                # completion.  Leave the loop.
+                break
+            else:
+                try:
+                    # Pass the generator's result along to whoever is
+                    # driving.  Capture the result as the next value to
+                    # send inward.
+                    value_in = yield value_out
+                except:
+                    # Or capture the exception if that's the flavor of the
+                    # next value.  This could possibly include GeneratorExit
+                    # which turns out to be just fine because throwing it into
+                    # the inner generator effectively propagates the close
+                    # (and with the right context!) just as you would want.
+                    # True, the GeneratorExit does get re-throwing out of the
+                    # gen.throw call and hits _the_generator_context's
+                    # contextmanager.  But @contextmanager extremely
+                    # conveniently eats it for us!  Thanks, @contextmanager!
+                    ok = False
+                    value_in = exc_info()
+                else:
+                    ok = True
+
+    wrapper.debug = False
+    return wrapper
diff --git a/eliot/twisted.py b/eliot/twisted.py
index 41cc0ed..9ea0932 100644
--- a/eliot/twisted.py
+++ b/eliot/twisted.py
@@ -9,11 +9,18 @@ import sys
 
 from twisted.logger import Logger as TwistedLogger
 from twisted.python.failure import Failure
+from twisted.internet.defer import inlineCallbacks
 
 from ._action import current_action
 from . import addDestination
+from ._generators import eliot_friendly_generator_function, use_generator_context
 
-__all__ = ["AlreadyFinished", "DeferredContext", "redirectLogsForTrial"]
+__all__ = [
+    "AlreadyFinished",
+    "DeferredContext",
+    "redirectLogsForTrial",
+    "inline_callbacks",
+]
 
 
 def _passthrough(result):
@@ -240,3 +247,17 @@ class _RedirectLogsForTrial(object):
 
 
 redirectLogsForTrial = _RedirectLogsForTrial(sys)
+
+
+def inline_callbacks(original, debug=False):
+    """
+    Decorate a function like ``inlineCallbacks`` would but in a more
+    Eliot-friendly way.  Use it just like ``inlineCallbacks`` but where you
+    want Eliot action contexts to Do The Right Thing inside the decorated
+    function.
+    """
+    use_generator_context()
+    f = eliot_friendly_generator_function(original)
+    if debug:
+        f.debug = True
+    return inlineCallbacks(f)


Test patch:
diff --git a/eliot/tests/test_generators.py b/eliot/tests/test_generators.py
new file mode 100644
index 0000000..ebd8d29
--- /dev/null
+++ b/eliot/tests/test_generators.py
@@ -0,0 +1,334 @@
+"""
+Tests for L{eliot._generators}.
+"""
+
+from __future__ import unicode_literals, absolute_import
+
+from pprint import pformat
+from unittest import TestCase
+
+from eliot import (
+    Message,
+    start_action,
+)
+from ..testing import (
+    capture_logging,
+    assertHasAction,
+)
+
+from .._generators import (
+    use_generator_context,
+    eliot_friendly_generator_function,
+)
+from .._action import _context_owner
+
+
+def assert_expected_action_tree(testcase, logger, expected_action_type, expected_type_tree):
+    """
+    Assert that a logger has a certain logged action with certain children.
+
+    @see: L{assert_generator_logs_action_tree}
+    """
+    logged_action = assertHasAction(
+        testcase,
+        logger,
+        expected_action_type,
+        True,
+    )
+    type_tree = logged_action.type_tree()
+    testcase.assertEqual(
+        {expected_action_type: expected_type_tree},
+        type_tree,
+        "Logger had messages:\n{}".format(pformat(logger.messages, indent=4)),
+    )
+
+
+def assert_generator_logs_action_tree(testcase, generator_function, logger, expected_action_type, expected_type_tree):
+    """
+    Assert that exhausting a generator from the given function logs an action
+    of the given type with children matching the given type tree.
+
+    @param testcase: A test case instance to use to make assertions.
+    @type testcase: L{unittest.TestCase}
+
+    @param generator_function: A no-argument callable that returns a generator
+        to be exhausted.
+
+    @param logger: A logger to inspect for logged messages.
+    @type logger: L{MemoryLogger}
+
+    @param expected_action_type: An action type which should be logged by the
+        generator.
+    @type expected_action_type: L{unicode}
+
+    @param expected_type_tree: The types of actions and messages which should
+        be logged beneath the expected action.  The structure of this value
+        matches the structure returned by L{LoggedAction.type_tree}.
+    @type expected_type_tree: L{list}
+    """
+    list(eliot_friendly_generator_function(generator_function)())
+    assert_expected_action_tree(
+        testcase,
+        logger,
+        expected_action_type,
+        expected_type_tree,
+    )
+
+
+class EliotFriendlyGeneratorFunctionTests(TestCase):
+    """
+    Tests for L{eliot_friendly_generator_function}.
+    """
+    # Get our custom assertion failure messages *and* the standard ones.
+    longMessage = True
+
+    def setUp(self):
+        use_generator_context()
+        self.addCleanup(_context_owner.reset)
+
+    @capture_logging(None)
+    def test_yield_none(self, logger):
+        @eliot_friendly_generator_function
+        def g():
+            Message.log(message_type=u"hello")
+            yield
+            Message.log(message_type=u"goodbye")
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type=u"the-action"):
+            list(g())
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action",
+            [u"hello", u"yielded", u"goodbye"],
+        )
+
+    @capture_logging(None)
+    def test_yield_value(self, logger):
+        expected = object()
+
+        @eliot_friendly_generator_function
+        def g():
+            Message.log(message_type=u"hello")
+            yield expected
+            Message.log(message_type=u"goodbye")
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type=u"the-action"):
+            self.assertEqual([expected], list(g()))
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action",
+            [u"hello", u"yielded", u"goodbye"],
+        )
+
+    @capture_logging(None)
+    def test_yield_inside_another_action(self, logger):
+        @eliot_friendly_generator_function
+        def g():
+            Message.log(message_type=u"a")
+            with start_action(action_type=u"confounding-factor"):
+                Message.log(message_type=u"b")
+                yield None
+                Message.log(message_type=u"c")
+            Message.log(message_type=u"d")
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type=u"the-action"):
+            list(g())
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action",
+            [u"a",
+             {u"confounding-factor": [u"b", u"yielded", u"c"]},
+             u"d",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_yield_inside_nested_actions(self, logger):
+        @eliot_friendly_generator_function
+        def g():
+            Message.log(message_type=u"a")
+            with start_action(action_type=u"confounding-factor"):
+                Message.log(message_type=u"b")
+                yield None
+                with start_action(action_type=u"double-confounding-factor"):
+                    yield None
+                    Message.log(message_type=u"c")
+                Message.log(message_type=u"d")
+            Message.log(message_type=u"e")
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type=u"the-action"):
+            list(g())
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                u"a",
+                {u"confounding-factor": [
+                    u"b",
+                    u"yielded",
+                    {u"double-confounding-factor": [
+                        u"yielded",
+                        u"c",
+                    ]},
+                    u"d",
+                ]},
+                u"e",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_generator_and_non_generator(self, logger):
+        @eliot_friendly_generator_function
+        def g():
+            Message.log(message_type=u"a")
+            yield
+            with start_action(action_type=u"action-a"):
+                Message.log(message_type=u"b")
+                yield
+                Message.log(message_type=u"c")
+
+            Message.log(message_type=u"d")
+            yield
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type=u"the-action"):
+            generator = g()
+            next(generator)
+            Message.log(message_type=u"0")
+            next(generator)
+            Message.log(message_type=u"1")
+            next(generator)
+            Message.log(message_type=u"2")
+            self.assertRaises(StopIteration, lambda: next(generator))
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                u"a",
+                u"yielded",
+                u"0",
+                {
+                    u"action-a": [
+                        u"b",
+                        u"yielded",
+                        u"c",
+                    ],
+                },
+                u"1",
+                u"d",
+                u"yielded",
+                u"2",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_concurrent_generators(self, logger):
+        @eliot_friendly_generator_function
+        def g(which):
+            Message.log(message_type=u"{}-a".format(which))
+            with start_action(action_type=which):
+                Message.log(message_type=u"{}-b".format(which))
+                yield
+                Message.log(message_type=u"{}-c".format(which))
+            Message.log(message_type=u"{}-d".format(which))
+        g.debug = True  # output yielded messages
+
+        gens = [g(u"1"), g(u"2")]
+        with start_action(action_type=u"the-action"):
+            while gens:
+                for g in gens[:]:
+                    try:
+                        next(g)
+                    except StopIteration:
+                        gens.remove(g)
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                u"1-a",
+                {u"1": [
+                    u"1-b",
+                    u"yielded",
+                    u"1-c",
+                ]},
+                u"2-a",
+                {u"2": [
+                    u"2-b",
+                    u"yielded",
+                    u"2-c",
+                ]},
+                u"1-d",
+                u"2-d",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_close_generator(self, logger):
+        @eliot_friendly_generator_function
+        def g():
+            Message.log(message_type=u"a")
+            try:
+                yield
+                Message.log(message_type=u"b")
+            finally:
+                Message.log(message_type=u"c")
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type=u"the-action"):
+            gen = g()
+            next(gen)
+            gen.close()
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                u"a",
+                u"yielded",
+                u"c",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_nested_generators(self, logger):
+        @eliot_friendly_generator_function
+        def g(recurse):
+            with start_action(action_type=u"a-recurse={}".format(recurse)):
+                Message.log(message_type=u"m-recurse={}".format(recurse))
+                if recurse:
+                    set(g(False))
+                else:
+                    yield
+
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type=u"the-action"):
+            set(g(True))
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [{
+                u"a-recurse=True": [
+                    u"m-recurse=True", {
+                        u"a-recurse=False": [
+                            u"m-recurse=False",
+                            u"yielded",
+                        ],
+                    },
+                ],
+            }],
+        )
diff --git a/eliot/tests/test_twisted.py b/eliot/tests/test_twisted.py
index d845104..256f1cd 100644
--- a/eliot/tests/test_twisted.py
+++ b/eliot/tests/test_twisted.py
@@ -8,7 +8,7 @@ import sys
 from functools import wraps
 
 try:
-    from twisted.internet.defer import Deferred, succeed, fail
+    from twisted.internet.defer import Deferred, succeed, fail, returnValue
     from twisted.trial.unittest import TestCase
     from twisted.python.failure import Failure
     from twisted.logger import globalLogPublisher
@@ -20,13 +20,18 @@ else:
     # logwriter.py causes a failure:
     from ..twisted import (
         DeferredContext, AlreadyFinished, _passthrough, redirectLogsForTrial,
-        _RedirectLogsForTrial, TwistedDestination
+        _RedirectLogsForTrial, TwistedDestination,
+        inline_callbacks,
     )
 
-from .._action import start_action, current_action, Action, TaskLevel
+from .test_generators import assert_expected_action_tree
+
+from .._action import (
+    start_action, current_action, Action, TaskLevel, _context_owner
+)
 from .._output import MemoryLogger, Logger
 from .._message import Message
-from ..testing import assertContainsFields
+from ..testing import assertContainsFields, capture_logging
 from .. import removeDestination, addDestination
 from .._traceback import write_traceback
 from .common import FakeSys
@@ -673,3 +678,207 @@ class TwistedDestinationTests(TestCase):
         self.assertEqual(
             writtenToTwisted, [("critical", written[0])]
         )
+
+
+class InlineCallbacksTests(TestCase):
+    """Tests for C{inline_callbacks}."""
+
+    # Get our custom assertion failure messages *and* the standard ones.
+    longMessage = True
+
+    def setUp(self):
+        # The generator sub-contxt is enabled automatically, so reset back to
+        # normal:
+        self.addCleanup(_context_owner.reset)
+
+    def _a_b_test(self, logger, g):
+        """A yield was done in between messages a and b inside C{inline_callbacks}."""
+        with start_action(action_type=u"the-action"):
+            self.assertIs(
+                None,
+                self.successResultOf(g()),
+            )
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                u"a",
+                u"yielded",
+                u"b",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_yield_none(self, logger):
+        def g():
+            Message.log(message_type=u"a")
+            yield
+            Message.log(message_type=u"b")
+        g = inline_callbacks(g, debug=True)
+
+        self._a_b_test(logger, g)
+
+    @capture_logging(None)
+    def test_yield_fired_deferred(self, logger):
+        def g():
+            Message.log(message_type=u"a")
+            yield succeed(None)
+            Message.log(message_type=u"b")
+        g = inline_callbacks(g, debug=True)
+
+        self._a_b_test(logger, g)
+
+    @capture_logging(None)
+    def test_yield_unfired_deferred(self, logger):
+        waiting = Deferred()
+
+        def g():
+            Message.log(message_type=u"a")
+            yield waiting
+            Message.log(message_type=u"b")
+        g = inline_callbacks(g, debug=True)
+
+        with start_action(action_type=u"the-action"):
+            d = g()
+            self.assertNoResult(waiting)
+            waiting.callback(None)
+            self.assertIs(
+                None,
+                self.successResultOf(d),
+            )
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                u"a",
+                u"yielded",
+                u"b",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_returnValue(self, logger):
+        result = object()
+
+        @inline_callbacks
+        def g():
+            if False:
+                yield
+            returnValue(result)
+
+        with start_action(action_type=u"the-action"):
+            d = g()
+            self.assertIs(
+                result,
+                self.successResultOf(d),
+            )
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action",
+            [],
+        )
+
+    @capture_logging(None)
+    def test_returnValue_in_action(self, logger):
+        result = object()
+
+        @inline_callbacks
+        def g():
+            if False:
+                yield
+            with start_action(action_type=u"g"):
+                returnValue(result)
+
+        with start_action(action_type=u"the-action"):
+            d = g()
+            self.assertIs(
+                result,
+                self.successResultOf(d),
+            )
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                {u"g": []},
+            ],
+        )
+
+    @capture_logging(None)
+    def test_nested_returnValue(self, logger):
+        result = object()
+        another = object()
+
+        def g():
+            d = h()
+            # Run h through to the end but ignore its result.
+            yield d
+            # Give back _our_ result.
+            returnValue(result)
+        g = inline_callbacks(g, debug=True)
+
+        def h():
+            yield
+            returnValue(another)
+        h = inline_callbacks(h, debug=True)
+
+        with start_action(action_type=u"the-action"):
+            d = g()
+            self.assertIs(
+                result,
+                self.successResultOf(d),
+            )
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            u"the-action", [
+                u"yielded",
+                u"yielded",
+            ],
+        )
+
+    @capture_logging(None)
+    def test_async_returnValue(self, logger):
+        result = object()
+        waiting = Deferred()
+
+        @inline_callbacks
+        def g():
+            yield waiting
+            returnValue(result)
+
+        with start_action(action_type=u"the-action"):
+            d = g()
+            waiting.callback(None)
+            self.assertIs(
+                result,
+                self.successResultOf(d),
+            )
+
+    @capture_logging(None)
+    def test_nested_async_returnValue(self, logger):
+        result = object()
+        another = object()
+
+        waiting = Deferred()
+
+        @inline_callbacks
+        def g():
+            yield h()
+            returnValue(result)
+
+        @inline_callbacks
+        def h():
+            yield waiting
+            returnValue(another)
+
+        with start_action(action_type=u"the-action"):
+            d = g()
+            waiting.callback(None)
+            self.assertIs(
+                result,
+                self.successResultOf(d),
+            )