From 10a8691ecade7ee3ebbf7223a9873bf83935f473 Mon Sep 17 00:00:00 2001
From: Jeffrey Rackauckas <jeffreyrack@gmail.com>
Date: Mon, 2 Jul 2018 19:46:26 -0700
Subject: [PATCH] Add support for yielded functions.

---
 src/_pytest/debugging.py | 47 ++++++++++++++++++++--------------------
 testing/test_pdb.py      | 21 ++++++++++++++++++
 2 files changed, 45 insertions(+), 23 deletions(-)

diff --git a/src/_pytest/debugging.py b/src/_pytest/debugging.py
index cd6656d18..2bf290de4 100644
--- a/src/_pytest/debugging.py
+++ b/src/_pytest/debugging.py
@@ -15,9 +15,6 @@ except ImportError:
     SUPPORTS_BREAKPOINT_BUILTIN = False
 
 
-immediately_break = False
-
-
 def pytest_addoption(parser):
     group = parser.getgroup("general")
     group._addoption(
@@ -49,8 +46,8 @@ def pytest_configure(config):
     else:
         pdb_cls = pdb.Pdb
 
-    global immediately_break
-    immediately_break = config.getvalue("trace")
+    if config.getvalue("trace"):
+        config.pluginmanager.register(PdbTrace(), "pdbtrace")
     if config.getvalue("usepdb"):
         config.pluginmanager.register(PdbInvoke(), "pdbinvoke")
 
@@ -76,24 +73,6 @@ def pytest_configure(config):
     config._cleanup.append(fin)
 
 
-@hookimpl(hookwrapper=True)
-def pytest_pyfunc_call(pyfuncitem):
-    if immediately_break:
-        pytestPDB.set_trace(set_break=False)
-        testfunction = pyfuncitem.obj
-        pyfuncitem.obj = pdb.runcall
-        if pyfuncitem._isyieldedfunction():
-            pyfuncitem._args = [testfunction, *pyfuncitem._args]
-        else:
-            if "func" in pyfuncitem._fixtureinfo.argnames:
-                raise ValueError("--trace can't be used with a fixture named func!")
-            pyfuncitem.funcargs["func"] = testfunction
-            new_list = list(pyfuncitem._fixtureinfo.argnames)
-            new_list.append("func")
-            pyfuncitem._fixtureinfo.argnames = tuple(new_list)
-    yield
-
-
 class pytestPDB(object):
     """ Pseudo PDB that defers to the real pdb. """
 
@@ -136,6 +115,28 @@ class PdbInvoke(object):
         post_mortem(tb)
 
 
+class PdbTrace(object):
+    @hookimpl(hookwrapper=True)
+    def pytest_pyfunc_call(self, pyfuncitem):
+        _test_pytest_function(pyfuncitem)
+        yield
+
+
+def _test_pytest_function(pyfuncitem):
+    pytestPDB.set_trace(set_break=False)
+    testfunction = pyfuncitem.obj
+    pyfuncitem.obj = pdb.runcall
+    if pyfuncitem._isyieldedfunction():
+        pyfuncitem._args = [testfunction, *pyfuncitem._args]
+    else:
+        if "func" in pyfuncitem._fixtureinfo.argnames:
+            raise ValueError("--trace can't be used with a fixture named func!")
+        pyfuncitem.funcargs["func"] = testfunction
+        new_list = list(pyfuncitem._fixtureinfo.argnames)
+        new_list.append("func")
+        pyfuncitem._fixtureinfo.argnames = tuple(new_list)
+
+
 def _enter_pdb(node, excinfo, rep):
     # XXX we re-use the TerminalReporter's terminalwriter
     # because this seems to avoid some encoding related troubles
diff --git a/testing/test_pdb.py b/testing/test_pdb.py
index 59aa183ad..316b033fe 100644
--- a/testing/test_pdb.py
+++ b/testing/test_pdb.py
@@ -714,3 +714,24 @@ class TestTraceOption:
         assert "1 passed" in rest
         assert "reading from stdin while output" not in rest
         TestPDB.flush(child)
+
+    def test_trace_against_yield_test(self, testdir):
+        p1 = testdir.makepyfile(
+            """
+            def is_equal(a, b):
+                assert a == b
+
+            def test_1():
+                assert is_equal, 1, 1
+            """
+        )
+                child = testdir.spawn_pytest("--trace " + str(p1))
+        child.expect("test_1")
+        child.expect("(Pdb)")
+        child.sendline("c")
+        child.expect("(Pdb)")
+        child.sendeof()
+        rest = child.read().decode("utf8")
+        assert "1 passed" in rest
+        assert "reading from stdin while output" not in rest
+        TestPDB.flush(child)