gh-150429: Fix sampling profiler generator stack test

Author: pablogsal

Lib/test/test_profiling/test_sampling_profiler/test_blocking.py

@@ -39,8 +39,9 @@ class TestBlockingModeStackAccuracy(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         # Test script that uses a generator consumed in a loop.
-        # When consume_generator is on the arithmetic lines (temp1, temp2, etc.),
-        # fibonacci_generator should NOT be in the stack at all.
+        # When consume_generator is the executing leaf frame on the arithmetic
+        # lines (temp1, temp2, etc.), fibonacci_generator should NOT be in the
+        # stack at all.
         # Line numbers are important here - see ARITHMETIC_LINES below.
         cls.generator_script = textwrap.dedent('''
             def fibonacci_generator(n):
@@ -65,29 +66,32 @@ def main():
             main()
         ''')
         # Line numbers of the arithmetic operations in consume_generator.
-        # These are the lines where fibonacci_generator should NOT be in the stack.
-        # The socket injection code adds 7 lines before our script.
-        # temp1 = value + 1  -> line 17
-        # temp2 = value * 2  -> line 18
-        # temp3 = value - 1  -> line 19
-        # result = ...       -> line 20
-        cls.ARITHMETIC_LINES = {17, 18, 19, 20}
+        # These are the lines where fibonacci_generator should NOT be in the
+        # stack when consume_generator is the executing leaf frame. They account
+        # for the socket prelude added by test_subprocess().
+        # temp1 = value + 1  -> line 16
+        # temp2 = value * 2  -> line 17
+        # temp3 = value - 1  -> line 18
+        # result = ...       -> line 19
+        cls.ARITHMETIC_LINES = {16, 17, 18, 19}
 
     def test_generator_not_under_consumer_arithmetic(self):
         """Test that fibonacci_generator doesn't appear when consume_generator does arithmetic.
 
-        When consume_generator is executing arithmetic lines (temp1, temp2, etc.),
-        fibonacci_generator should NOT be anywhere in the stack - it's not being
-        called at that point.
+        When consume_generator is the leaf frame on arithmetic lines (temp1,
+        temp2, etc.), fibonacci_generator should NOT be anywhere in the stack -
+        it's not being called at that point. Non-leaf frame line numbers are
+        caller/resume metadata, not proof that the frame is executing.
 
         Valid stacks:
-          - consume_generator at 'for value in gen:' line WITH fibonacci_generator
-            at the top (generator is yielding)
+          - fibonacci_generator at the top (generator is executing), with
+            consume_generator below it
           - consume_generator at arithmetic lines WITHOUT fibonacci_generator
             (we're just doing math, not calling the generator)
 
         Invalid stacks (indicate torn/inconsistent reads):
-          - consume_generator at arithmetic lines WITH fibonacci_generator
+          - consume_generator leaf frame at arithmetic lines WITH
+            fibonacci_generator
             anywhere in the stack
 
         Note: call_tree is ordered from bottom (index 0) to top (index -1).
@@ -110,22 +114,30 @@ def test_generator_not_under_consumer_arithmetic(self):
         total_samples = 0
         invalid_stacks = 0
         arithmetic_samples = 0
+        generator_samples = 0
+        generator_not_leaf_samples = 0
 
         for (call_tree, _thread_id), count in collector.stack_counter.items():
             total_samples += count
 
             if not call_tree:
                 continue
 
-            # Find consume_generator in the stack and check its line number
-            for i, (filename, lineno, funcname) in enumerate(call_tree):
-                if funcname == "consume_generator" and lineno in self.ARITHMETIC_LINES:
-                    arithmetic_samples += count
-                    # Check if fibonacci_generator appears anywhere in this stack
-                    func_names = [frame[2] for frame in call_tree]
-                    if "fibonacci_generator" in func_names:
-                        invalid_stacks += count
-                    break
+            # Non-leaf frame line numbers can point at resume locations while
+            # a callee is the executing leaf frame.
+            _, lineno, funcname = call_tree[-1]
+            func_names = [frame[2] for frame in call_tree]
+
+            if "fibonacci_generator" in func_names:
+                generator_samples += count
+                if funcname != "fibonacci_generator":
+                    generator_not_leaf_samples += count
+
+            if funcname == "consume_generator" and lineno in self.ARITHMETIC_LINES:
+                arithmetic_samples += count
+                # Check if fibonacci_generator appears anywhere in this stack.
+                if "fibonacci_generator" in func_names:
+                    invalid_stacks += count
 
         self.assertGreater(total_samples, 10,
             f"Expected at least 10 samples, got {total_samples}")
@@ -134,8 +146,15 @@ def test_generator_not_under_consumer_arithmetic(self):
         self.assertGreater(arithmetic_samples, 0,
             f"Expected some samples on arithmetic lines, got {arithmetic_samples}")
 
+        self.assertGreater(generator_samples, 0,
+            f"Expected some samples in fibonacci_generator, got {generator_samples}")
+
+        self.assertEqual(generator_not_leaf_samples, 0,
+            f"Found {generator_not_leaf_samples}/{generator_samples} stacks where "
+            f"fibonacci_generator appears but is not the leaf frame.")
+
         self.assertEqual(invalid_stacks, 0,
             f"Found {invalid_stacks}/{arithmetic_samples} invalid stacks where "
             f"fibonacci_generator appears in the stack when consume_generator "
-            f"is on an arithmetic line. This indicates torn/inconsistent stack "
-            f"traces are being captured.")
+            f"is the leaf frame on an arithmetic line. This indicates "
+            f"torn/inconsistent stack traces are being captured.")