Fix tests

Author: pablogsal

Lib/profiling/sampling/cli.py

@@ -390,6 +390,19 @@ def _validate_args(args, parser):
             "Live mode requires the curses module, which is not available."
         )
 
+    # Async-aware mode is incompatible with --native and --gc/--no-gc
+    if args.async_aware is not None:
+        issues = []
+        if args.native:
+            issues.append("--native")
+        if not args.gc:
+            issues.append("--no-gc")
+        if issues:
+            parser.error(
+                f"Options {', '.join(issues)} are incompatible with --async-aware. "
+                "Async-aware profiling uses task-based stack reconstruction."
+            )
+
     # Live mode is incompatible with format options
     if hasattr(args, 'live') and args.live:
         if args.format != "pstats":

Lib/profiling/sampling/live_collector/collector.py

@@ -308,26 +308,6 @@ def _get_async_frame_iterator(self, stack_frames):
         for frames, thread_id, task_id in self._iter_async_frames(stack_frames):
             yield frames, thread_id
 
-    def _collect_sync_thread_stats(self, stack_frames, has_gc_frame):
-        """Collect thread status stats for sync mode."""
-        status_counts, sample_has_gc, per_thread_stats = self._collect_thread_status_stats(stack_frames)
-        for key, count in status_counts.items():
-            self.thread_status_counts[key] += count
-        if sample_has_gc:
-            has_gc_frame = True
-
-        for thread_id, stats in per_thread_stats.items():
-            thread_data = self._get_or_create_thread_data(thread_id)
-            thread_data.has_gil += stats.get("has_gil", 0)
-            thread_data.on_cpu += stats.get("on_cpu", 0)
-            thread_data.gil_requested += stats.get("gil_requested", 0)
-            thread_data.unknown += stats.get("unknown", 0)
-            thread_data.total += stats.get("total", 0)
-            if stats.get("gc_samples", 0):
-                thread_data.gc_frame_samples += stats["gc_samples"]
-
-        return has_gc_frame
-
     def collect_failed_sample(self):
         self.failed_samples += 1
         self.total_samples += 1
@@ -342,7 +322,21 @@ def collect(self, stack_frames):
 
         # Collect thread status stats (only available in sync mode)
         if not self.async_aware:
-            has_gc_frame = self._collect_sync_thread_stats(stack_frames, has_gc_frame)
+            status_counts, sample_has_gc, per_thread_stats = self._collect_thread_status_stats(stack_frames)
+            for key, count in status_counts.items():
+                self.thread_status_counts[key] += count
+            if sample_has_gc:
+                has_gc_frame = True
+
+            for thread_id, stats in per_thread_stats.items():
+                thread_data = self._get_or_create_thread_data(thread_id)
+                thread_data.has_gil += stats.get("has_gil", 0)
+                thread_data.on_cpu += stats.get("on_cpu", 0)
+                thread_data.gil_requested += stats.get("gil_requested", 0)
+                thread_data.unknown += stats.get("unknown", 0)
+                thread_data.total += stats.get("total", 0)
+                if stats.get("gc_samples", 0):
+                    thread_data.gc_frame_samples += stats["gc_samples"]
 
         # Process frames using pre-selected iterator
         for frames, thread_id in self._get_frame_iterator(stack_frames):
@@ -355,20 +349,9 @@ def collect(self, stack_frames):
             if thread_id is not None and thread_id not in self.thread_ids:
                 self.thread_ids.append(thread_id)
 
-            # Check for GC frames and update per-thread sample count
-            thread_has_gc_frame = False
-            for frame in frames:
-                funcname = getattr(frame, "funcname", "")
-                if "<GC>" in funcname or "gc_collect" in funcname:
-                    has_gc_frame = True
-                    thread_has_gc_frame = True
-                    break
-
             if thread_id is not None:
                 thread_data = self._get_or_create_thread_data(thread_id)
                 thread_data.sample_count += 1
-                if thread_has_gc_frame:
-                    thread_data.gc_frame_samples += 1
 
         if has_gc_frame:
             self.gc_frame_samples += 1

Lib/test/test_profiling/test_sampling_profiler/test_async.py

@@ -697,79 +697,102 @@ def test_sample_profiler_sample_accepts_async_aware(self):
         sig = inspect.signature(SampleProfiler.sample)
         self.assertIn("async_aware", sig.parameters)
 
-    def test_async_aware_all_uses_get_all_awaited_by(self):
-        """Test that async_aware='all' calls get_all_awaited_by on unwinder."""
-        from unittest.mock import Mock, patch
-        from profiling.sampling.sample import SampleProfiler
-
-        with patch('profiling.sampling.sample._remote_debugging') as mock_rd:
-            mock_unwinder = Mock()
-            mock_unwinder.get_all_awaited_by.return_value = []
-            mock_rd.RemoteUnwinder.return_value = mock_unwinder
-
-            profiler = SampleProfiler(
-                pid=12345,
-                sample_interval_usec=1000,
-                all_threads=False
-            )
-            profiler.unwinder = mock_unwinder
-
-            mock_collector = Mock()
-            mock_collector.running = False  # Stop immediately
+    def test_async_aware_all_sees_sleeping_and_running_tasks(self):
+        """Test async_aware='all' captures both sleeping and CPU-running tasks."""
+        # Sleeping task (awaiting)
+        sleeping_task = MockTaskInfo(
+            task_id=1,
+            task_name="SleepingTask",
+            coroutine_stack=[
+                MockCoroInfo(
+                    task_name="SleepingTask",
+                    call_stack=[MockFrameInfo("sleeper.py", 10, "sleep_work")]
+                )
+            ],
+            awaited_by=[]
+        )
 
-            # Sample with async_aware="all"
-            profiler.sample(mock_collector, duration_sec=0.001, async_aware="all")
+        # CPU-running task (active)
+        running_task = MockTaskInfo(
+            task_id=2,
+            task_name="RunningTask",
+            coroutine_stack=[
+                MockCoroInfo(
+                    task_name="RunningTask",
+                    call_stack=[MockFrameInfo("runner.py", 20, "cpu_work")]
+                )
+            ],
+            awaited_by=[]
+        )
 
-            # Should have called get_all_awaited_by
-            mock_unwinder.get_all_awaited_by.assert_called()
+        # Both tasks returned by get_all_awaited_by
+        awaited_info_list = [MockAwaitedInfo(thread_id=100, awaited_by=[sleeping_task, running_task])]
 
-    def test_async_aware_running_uses_get_async_stack_trace(self):
-        """Test that async_aware='running' calls get_async_stack_trace on unwinder."""
-        from unittest.mock import Mock, patch
-        from profiling.sampling.sample import SampleProfiler
+        collector = PstatsCollector(sample_interval_usec=1000)
+        collector.collect(awaited_info_list)
+        collector.create_stats()
 
-        with patch('profiling.sampling.sample._remote_debugging') as mock_rd:
-            mock_unwinder = Mock()
-            mock_unwinder.get_async_stack_trace.return_value = []
-            mock_rd.RemoteUnwinder.return_value = mock_unwinder
+        # Both tasks should be visible
+        sleeping_key = ("sleeper.py", 10, "sleep_work")
+        running_key = ("runner.py", 20, "cpu_work")
 
-            profiler = SampleProfiler(
-                pid=12345,
-                sample_interval_usec=1000,
-                all_threads=False
-            )
-            profiler.unwinder = mock_unwinder
+        self.assertIn(sleeping_key, collector.stats)
+        self.assertIn(running_key, collector.stats)
 
-            mock_collector = Mock()
-            mock_collector.running = False
+        # Task markers should also be present
+        task_keys = [k for k in collector.stats if k[0] == "<task>"]
+        self.assertGreater(len(task_keys), 0, "Should have <task> markers in stats")
 
-            profiler.sample(mock_collector, duration_sec=0.001, async_aware="running")
+        # Verify task names are in the markers
+        task_names = [k[2] for k in task_keys]
+        self.assertTrue(
+            any("SleepingTask" in name for name in task_names),
+            "SleepingTask should be in task markers"
+        )
+        self.assertTrue(
+            any("RunningTask" in name for name in task_names),
+            "RunningTask should be in task markers"
+        )
 
-            mock_unwinder.get_async_stack_trace.assert_called()
+    def test_async_aware_running_sees_only_running_task(self):
+        """Test async_aware='running' only shows the currently running task stack."""
+        # Only the running task's stack is returned by get_async_stack_trace
+        running_task = MockTaskInfo(
+            task_id=2,
+            task_name="RunningTask",
+            coroutine_stack=[
+                MockCoroInfo(
+                    task_name="RunningTask",
+                    call_stack=[MockFrameInfo("runner.py", 20, "cpu_work")]
+                )
+            ],
+            awaited_by=[]
+        )
 
-    def test_async_aware_none_uses_get_stack_trace(self):
-        """Test that async_aware=None uses regular get_stack_trace."""
-        from unittest.mock import Mock, patch
-        from profiling.sampling.sample import SampleProfiler
+        # get_async_stack_trace only returns the running task
+        awaited_info_list = [MockAwaitedInfo(thread_id=100, awaited_by=[running_task])]
 
-        with patch('profiling.sampling.sample._remote_debugging') as mock_rd:
-            mock_unwinder = Mock()
-            mock_unwinder.get_stack_trace.return_value = []
-            mock_rd.RemoteUnwinder.return_value = mock_unwinder
+        collector = PstatsCollector(sample_interval_usec=1000)
+        collector.collect(awaited_info_list)
+        collector.create_stats()
 
-            profiler = SampleProfiler(
-                pid=12345,
-                sample_interval_usec=1000,
-                all_threads=False
-            )
-            profiler.unwinder = mock_unwinder
+        # Only running task should be visible
+        running_key = ("runner.py", 20, "cpu_work")
+        self.assertIn(running_key, collector.stats)
 
-            mock_collector = Mock()
-            mock_collector.running = False
+        # Verify we don't see the sleeping task (it wasn't in the input)
+        sleeping_key = ("sleeper.py", 10, "sleep_work")
+        self.assertNotIn(sleeping_key, collector.stats)
 
-            profiler.sample(mock_collector, duration_sec=0.001, async_aware=None)
+        # Task marker for running task should be present
+        task_keys = [k for k in collector.stats if k[0] == "<task>"]
+        self.assertGreater(len(task_keys), 0, "Should have <task> markers in stats")
 
-            mock_unwinder.get_stack_trace.assert_called()
+        task_names = [k[2] for k in task_keys]
+        self.assertTrue(
+            any("RunningTask" in name for name in task_names),
+            "RunningTask should be in task markers"
+        )
 
 
 if __name__ == "__main__":

Lib/test/test_profiling/test_sampling_profiler/test_cli.py

@@ -607,3 +607,55 @@ def test_async_aware_invalid_choice(self):
             main()
 
         self.assertEqual(cm.exception.code, 2)  # argparse error
+
+    def test_async_aware_incompatible_with_native(self):
+        """Test --async-aware is incompatible with --native."""
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "all", "--native"]
+
+        with (
+            mock.patch("sys.argv", test_args),
+            mock.patch("sys.stderr", io.StringIO()) as mock_stderr,
+            self.assertRaises(SystemExit) as cm,
+        ):
+            from profiling.sampling.cli import main
+            main()
+
+        self.assertEqual(cm.exception.code, 2)  # argparse error
+        error_msg = mock_stderr.getvalue()
+        self.assertIn("--native", error_msg)
+        self.assertIn("incompatible with --async-aware", error_msg)
+
+    def test_async_aware_incompatible_with_no_gc(self):
+        """Test --async-aware is incompatible with --no-gc."""
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "running", "--no-gc"]
+
+        with (
+            mock.patch("sys.argv", test_args),
+            mock.patch("sys.stderr", io.StringIO()) as mock_stderr,
+            self.assertRaises(SystemExit) as cm,
+        ):
+            from profiling.sampling.cli import main
+            main()
+
+        self.assertEqual(cm.exception.code, 2)  # argparse error
+        error_msg = mock_stderr.getvalue()
+        self.assertIn("--no-gc", error_msg)
+        self.assertIn("incompatible with --async-aware", error_msg)
+
+    def test_async_aware_incompatible_with_both_native_and_no_gc(self):
+        """Test --async-aware is incompatible with both --native and --no-gc."""
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "all", "--native", "--no-gc"]
+
+        with (
+            mock.patch("sys.argv", test_args),
+            mock.patch("sys.stderr", io.StringIO()) as mock_stderr,
+            self.assertRaises(SystemExit) as cm,
+        ):
+            from profiling.sampling.cli import main
+            main()
+
+        self.assertEqual(cm.exception.code, 2)  # argparse error
+        error_msg = mock_stderr.getvalue()
+        self.assertIn("--native", error_msg)
+        self.assertIn("--no-gc", error_msg)
+        self.assertIn("incompatible with --async-aware", error_msg)