1 files changed, 108 insertions, 12 deletions
diff --git a/src/greenlet/tests/test_tracing.py b/src/greenlet/tests/test_tracing.py
index ce059ce..2ab4d71 100644
--- a/src/greenlet/tests/test_tracing.py
+++ b/src/greenlet/tests/test_tracing.py
@@ -111,7 +111,7 @@ class TestPythonTracing(unittest.TestCase):
             ('call', 'tpt_callback'),
             ('return', 'tpt_callback'),
             ('call', '__exit__'),
-            ('c_call', '__exit__')
+            ('c_call', '__exit__'),
         ])
 
     def _trace_switch(self, glet):
@@ -130,7 +130,7 @@ class TestPythonTracing(unittest.TestCase):
             ('return', 'run'),
             ('c_return', '_trace_switch'),
             ('call', '__exit__'),
-            ('c_call', '__exit__')
+            ('c_call', '__exit__'),
         ])
 
     def test_trace_events_into_greenlet_func_already_set(self):
@@ -145,27 +145,123 @@ class TestPythonTracing(unittest.TestCase):
                 return tpt_callback()
         self._check_trace_events_func_already_set(X())
 
-    def test_trace_events_from_greenlet_sets_func(self):
+    def _check_trace_events_from_greenlet_sets_profiler(self, g, tracer):
+        g.switch()
+        tpt_callback()
+        tracer.__exit__()
+        self.assertEqual(tracer.actions, [
+            ('return', '__enter__'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('return', 'run'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('call', '__exit__'),
+            ('c_call', '__exit__'),
+        ])
+
+
+    def test_trace_events_from_greenlet_func_sets_profiler(self):
         tracer = PythonTracer()
         def run():
             tracer.__enter__()
             return tpt_callback()
 
-        g = greenlet.greenlet(run)
-        g.switch()
-        tpt_callback()
-        tracer.__exit__()
+        self._check_trace_events_from_greenlet_sets_profiler(greenlet.greenlet(run),
+                                                             tracer)
+
+    def test_trace_events_from_greenlet_subclass_sets_profiler(self):
+        tracer = PythonTracer()
+        class X(greenlet.greenlet):
+            def run(self):
+                tracer.__enter__()
+                return tpt_callback()
+
+        self._check_trace_events_from_greenlet_sets_profiler(X(), tracer)
+
+
+    def test_trace_events_multiple_greenlets_switching(self):
+        tracer = PythonTracer()
+
+        g1 = None
+        g2 = None
+
+        def g1_run():
+            tracer.__enter__()
+            tpt_callback()
+            g2.switch()
+            tpt_callback()
+            return 42
+
+        def g2_run():
+            tpt_callback()
+            tracer.__exit__()
+            tpt_callback()
+            g1.switch()
+
+        g1 = greenlet.greenlet(g1_run)
+        g2 = greenlet.greenlet(g2_run)
+
+        x = g1.switch()
+        self.assertEqual(x, 42)
+        tpt_callback() # ensure not in the trace
         self.assertEqual(tracer.actions, [
             ('return', '__enter__'),
             ('call', 'tpt_callback'),
             ('return', 'tpt_callback'),
-            ('return', 'run'),
+            ('c_call', 'g1_run'),
+            ('call', 'g2_run'),
             ('call', 'tpt_callback'),
             ('return', 'tpt_callback'),
             ('call', '__exit__'),
-            ('c_call', '__exit__')
+            ('c_call', '__exit__'),
         ])
 
-    # What about multiple greenlets switching amongst each other? Probably doesn't
-    # work either.
-    # What about with throw()? And throw() as the initial action?
+    def test_trace_events_multiple_greenlets_switching_siblings(self):
+        # Like the first version, but get both greenlets running first
+        # as "siblings" and then establish the tracing.
+        tracer = PythonTracer()
+
+        g1 = None
+        g2 = None
+
+        def g1_run():
+            greenlet.getcurrent().parent.switch()
+            tracer.__enter__()
+            tpt_callback()
+            g2.switch()
+            tpt_callback()
+            return 42
+
+        def g2_run():
+            greenlet.getcurrent().parent.switch()
+
+            tpt_callback()
+            tracer.__exit__()
+            tpt_callback()
+            g1.switch()
+
+        g1 = greenlet.greenlet(g1_run)
+        g2 = greenlet.greenlet(g2_run)
+
+        # Start g1
+        g1.switch()
+        # And it immediately returns control to us.
+        # Start g2
+        g2.switch()
+        # Which also returns. Now kick of the real part of the
+        # test.
+        x = g1.switch()
+        self.assertEqual(x, 42)
+
+        tpt_callback() # ensure not in the trace
+        self.assertEqual(tracer.actions, [
+            ('return', '__enter__'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('c_call', 'g1_run'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('call', '__exit__'),
+            ('c_call', '__exit__'),
+        ])