From 9f5b2019f878a99e93c0bafda0c85b3f12caf301 Mon Sep 17 00:00:00 2001
From: Tom Prince <tom.prince@private.storage>
Date: Wed, 14 Jul 2021 16:43:13 -0600
Subject: [PATCH] Capture contect from first call in
 `eliot_friendly_generator_function`.

---
 eliot/_generators.py           | 17 ++++++++----
 eliot/tests/test_generators.py | 50 +++++++++++++++++++++++++++++-----
 2 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/eliot/_generators.py b/eliot/_generators.py
index f398b4b..bb452ef 100644
--- a/eliot/_generators.py
+++ b/eliot/_generators.py
@@ -43,7 +43,7 @@ def eliot_friendly_generator_function(original):
     """
 
     @wraps(original)
-    def wrapper(*a, **kw):
+    def wrapper(context, *a, **kw):
         # Keep track of whether the next value to deliver to the generator is
         # a non-exception or an exception.
         ok = True
@@ -57,8 +57,6 @@ def wrapper(*a, **kw):
         # generator function can run until we call send or throw on it.
         gen = original(*a, **kw)
 
-        # Initialize the per-generator context to a copy of the current context.
-        context = copy_context()
         while True:
             try:
                 # Whichever way we invoke the generator, we will do it
@@ -97,7 +95,7 @@ def go():
                     # indication of where the yield occurred.
                     #
                     # This is noisy, enable only for debugging:
-                    if wrapper.debug:
+                    if trampoline.debug:
                         log_message(message_type="yielded")
                     return value_out
 
@@ -127,5 +125,12 @@ def go():
                 else:
                     ok = True
 
-    wrapper.debug = False
-    return wrapper
+    @wraps(original)
+    def trampoline(*a, **kw):
+        # Initialize the generator context to a copy of the current context at
+        # the site where the generator is invoked.
+        context = copy_context()
+        return wrapper(context, *a, **kw)
+
+    trampoline.debug = False
+    return trampoline
diff --git a/eliot/tests/test_generators.py b/eliot/tests/test_generators.py
index a92cff3..7a216e9 100644
--- a/eliot/tests/test_generators.py
+++ b/eliot/tests/test_generators.py
@@ -220,14 +220,14 @@ def g(which):
 
         g.debug = True  # output yielded messages
 
-        gens = [g("1"), g("2")]
         with start_action(action_type="the-action"):
-            while gens:
-                for g in gens[:]:
-                    try:
-                        next(g)
-                    except StopIteration:
-                        gens.remove(g)
+            gens = [g("1"), g("2")]
+        while gens:
+            for g in gens[:]:
+                try:
+                    next(g)
+                except StopIteration:
+                    gens.remove(g)
 
         assert_expected_action_tree(
             self,
@@ -292,3 +292,39 @@ def g(recurse):
                 }
             ],
         )
+
+    @capture_logging(None)
+    def test_capture_context(self, logger):
+        """
+        L{eliot_friendly_generator_function} decorated generators capture the
+        context where they are created, not where L{.send} is first called on
+        them.
+        """
+
+        @eliot_friendly_generator_function
+        def g():
+            yield
+
+        g.debug = True  # output yielded messages
+
+        with start_action(action_type="the-action"):
+            with start_action(action_type="start"):
+                gen = g()
+            with start_action(action_type="run"):
+                list(gen)
+
+        assert_expected_action_tree(
+            self,
+            logger,
+            "the-action",
+            [
+                {
+                    "start": [
+                        "yielded",
+                    ],
+                },
+                {
+                    "run": [],
+                },
+            ],
+        )