fixup! test

Author: pablogsal

Lib/test/test_profiling/test_sampling_profiler/test_collectors.py

@@ -11,6 +11,7 @@
 
 try:
     import _remote_debugging  # noqa: F401
+    from profiling.sampling import gecko_collector
     from profiling.sampling.pstats_collector import PstatsCollector
     from profiling.sampling.stack_collector import (
         CollapsedStackCollector,
@@ -59,6 +60,42 @@ def find_child_by_name(children, strings, substr):
     return None
 
 
+def export_gecko_profile(testcase, collector):
+    gecko_out = tempfile.NamedTemporaryFile(suffix=".json", delete=False)
+    testcase.addCleanup(close_and_unlink, gecko_out)
+    # We cannot overwrite an open file on Windows.
+    gecko_out.close()
+
+    with captured_stdout(), captured_stderr():
+        collector.export(gecko_out.name)
+
+    testcase.assertGreater(os.path.getsize(gecko_out.name), 0)
+    with open(gecko_out.name, encoding="utf-8") as file:
+        return json.load(file)
+
+
+def assert_gecko_column_lengths(testcase, table, columns):
+    expected = table["length"]
+    for column in columns:
+        testcase.assertEqual(
+            len(table[column]), expected,
+            f"{column!r} has wrong length",
+        )
+
+
+def gecko_marker_names(profile, markers):
+    string_array = profile["shared"]["stringArray"]
+    return [string_array[idx] for idx in markers["name"]]
+
+
+def gecko_opcode_marker_data(profile):
+    markers = profile["threads"][0]["markers"]
+    return [
+        data for data in markers["data"]
+        if data.get("type") == "Opcode"
+    ]
+
+
 class TestSampleProfilerComponents(unittest.TestCase):
     """Unit tests for individual profiler components."""
 
@@ -583,9 +620,10 @@ def test_gecko_collector_basic(self):
 
         # Verify samples
         samples = thread_data["samples"]
-        self.assertEqual(len(samples["stack"]), 1)
-        self.assertEqual(len(samples["time"]), 1)
         self.assertEqual(samples["length"], 1)
+        assert_gecko_column_lengths(
+            self, samples, ("stack", "time", "eventDelay")
+        )
 
         # Verify function table structure and content
         func_table = thread_data["funcTable"]
@@ -622,11 +660,6 @@ def test_gecko_collector_basic(self):
     @unittest.skipIf(is_emscripten, "threads not available")
     def test_gecko_collector_export(self):
         """Test Gecko profile export functionality."""
-        gecko_out = tempfile.NamedTemporaryFile(suffix=".json", delete=False)
-        self.addCleanup(close_and_unlink, gecko_out)
-        # We cannot overwrite an open file on Windows.
-        gecko_out.close()
-
         collector = GeckoCollector(1000)
 
         test_frames1 = [
@@ -659,17 +692,7 @@ def test_gecko_collector_export(self):
         collector.collect(test_frames2)
         collector.collect(test_frames3)
 
-        # Export gecko profile
-        with captured_stdout(), captured_stderr():
-            collector.export(gecko_out.name)
-
-        # Verify file was created and contains valid data
-        self.assertTrue(os.path.exists(gecko_out.name))
-        self.assertGreater(os.path.getsize(gecko_out.name), 0)
-
-        # Check file contains valid JSON
-        with open(gecko_out.name, "r") as f:
-            profile_data = json.load(f)
+        profile_data = export_gecko_profile(self, collector)
 
         # Should be valid Gecko profile format
         self.assertIn("meta", profile_data)
@@ -690,6 +713,100 @@ def test_gecko_collector_export(self):
         self.assertIn("func2", string_array)
         self.assertIn("other_func", string_array)
 
+        thread_data = profile_data["threads"][0]
+        assert_gecko_column_lengths(
+            self, thread_data["samples"], ("stack", "time", "eventDelay")
+        )
+
+    @unittest.skipIf(is_emscripten, "threads not available")
+    def test_gecko_collector_export_after_spill_flush(self):
+        """Test Gecko profile export after spill buffers flush to disk."""
+        old_buffer_bytes = gecko_collector.DEFAULT_SPILL_BUFFER_BYTES
+        gecko_collector.DEFAULT_SPILL_BUFFER_BYTES = 1
+        self.addCleanup(
+            setattr, gecko_collector, "DEFAULT_SPILL_BUFFER_BYTES",
+            old_buffer_bytes
+        )
+
+        collector = GeckoCollector(1000)
+        test_frames = [
+            MockInterpreterInfo(
+                0,
+                [
+                    MockThreadInfo(
+                        1,
+                        [MockFrameInfo("file.py", 10, "func")],
+                        status=THREAD_STATUS_HAS_GIL,
+                    )
+                ],
+            )
+        ]
+        collector.collect(test_frames, timestamps_us=[1000, 2000, 3000])
+
+        profile_data = export_gecko_profile(self, collector)
+        samples = profile_data["threads"][0]["samples"]
+        self.assertEqual(samples["length"], 3)
+        assert_gecko_column_lengths(
+            self, samples, ("stack", "time", "eventDelay")
+        )
+
+    @unittest.skipIf(is_emscripten, "threads not available")
+    def test_gecko_collector_rejects_collect_after_export(self):
+        collector = GeckoCollector(1000)
+        test_frames = [
+            MockInterpreterInfo(
+                0,
+                [
+                    MockThreadInfo(
+                        1,
+                        [MockFrameInfo("file.py", 10, "func")],
+                        status=THREAD_STATUS_HAS_GIL,
+                    )
+                ],
+            )
+        ]
+        collector.collect(test_frames)
+        export_gecko_profile(self, collector)
+
+        with self.assertRaisesRegex(RuntimeError, "after export"):
+            collector.collect(test_frames)
+
+    @unittest.skipIf(is_emscripten, "threads not available")
+    def test_gecko_collector_export_failure_keeps_existing_file(self):
+        collector = GeckoCollector(1000)
+        test_frames = [
+            MockInterpreterInfo(
+                0,
+                [
+                    MockThreadInfo(
+                        1,
+                        [MockFrameInfo("file.py", 10, "func")],
+                        status=THREAD_STATUS_HAS_GIL,
+                    )
+                ],
+            )
+        ]
+        collector.collect(test_frames)
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            filename = os.path.join(temp_dir, "profile.json")
+            with open(filename, "w", encoding="utf-8") as file:
+                file.write("existing")
+
+            before = set(os.listdir(temp_dir))
+
+            def fail(file):
+                raise OSError("boom")
+
+            collector._stream_profile = fail
+            with captured_stdout(), captured_stderr():
+                with self.assertRaisesRegex(OSError, "boom"):
+                    collector.export(filename)
+
+            with open(filename, encoding="utf-8") as file:
+                self.assertEqual(file.read(), "existing")
+            self.assertEqual(set(os.listdir(temp_dir)), before)
+
     def test_gecko_collector_markers(self):
         """Test Gecko profile markers for GIL and CPU state tracking."""
         collector = GeckoCollector(1000)
@@ -773,21 +890,16 @@ def test_gecko_collector_markers(self):
         self.assertIn("markers", thread_data)
         markers = thread_data["markers"]
 
-        # Should have marker arrays
-        self.assertIn("name", markers)
-        self.assertIn("startTime", markers)
-        self.assertIn("endTime", markers)
-        self.assertIn("category", markers)
         self.assertGreater(
             markers["length"], 0, "Should have generated markers"
         )
-
-        # Get marker names from string table
-        string_array = profile_data["shared"]["stringArray"]
-        marker_names = [string_array[idx] for idx in markers["name"]]
+        assert_gecko_column_lengths(
+            self, markers,
+            ("data", "name", "startTime", "endTime", "phase", "category"),
+        )
 
         # Verify we have different marker types
-        marker_name_set = set(marker_names)
+        marker_name_set = set(gecko_marker_names(profile_data, markers))
 
         # Should have "Has GIL" markers (when thread had GIL)
         self.assertIn(
@@ -2661,7 +2773,7 @@ def test_gecko_collector_opcodes_enabled(self):
     def test_gecko_opcode_state_tracking(self):
         """Test that GeckoCollector tracks opcode state changes."""
         collector = GeckoCollector(sample_interval_usec=1000, opcodes=True)
-        self.addCleanup(lambda: collector.spill_dir.cleanup())
+        self.addCleanup(collector._cleanup_spills)
 
         # First sample with opcode 90 (RAISE_VARARGS)
         frame1 = MockFrameInfo("test.py", 10, "func", opcode=90)
@@ -2683,7 +2795,6 @@ def test_gecko_opcode_state_tracking(self):
     def test_gecko_opcode_state_change_emits_marker(self):
         """Test that opcode state change emits an interval marker."""
         collector = GeckoCollector(sample_interval_usec=1000, opcodes=True)
-        self.addCleanup(lambda: collector.spill_dir.cleanup())
 
         # First sample: opcode 90
         frame1 = MockFrameInfo("test.py", 10, "func", opcode=90)
@@ -2706,12 +2817,32 @@ def test_gecko_opcode_state_change_emits_marker(self):
         collector.collect(frames2)
 
         # Should have emitted a marker for the first opcode
-        self.assertGreater(collector.thread_spills[1].marker_count, 0)
+        profile = collector._build_profile()
+        markers = profile["threads"][0]["markers"]
+        assert_gecko_column_lengths(
+            self, markers,
+            ("data", "name", "startTime", "endTime", "phase", "category"),
+        )
+        opcode_markers = gecko_opcode_marker_data(profile)
+        self.assertIn(
+            {
+                "opcode": 90,
+                "line": 10,
+                "function": "func",
+            },
+            [
+                {
+                    "opcode": marker["opcode"],
+                    "line": marker["line"],
+                    "function": marker["function"],
+                }
+                for marker in opcode_markers
+            ],
+        )
 
     def test_gecko_opcode_markers_not_emitted_when_disabled(self):
         """Test that no opcode markers when opcodes=False."""
         collector = GeckoCollector(sample_interval_usec=1000, opcodes=False)
-        self.addCleanup(lambda: collector.spill_dir.cleanup())
 
         frame1 = MockFrameInfo("test.py", 10, "func", opcode=90)
         frames1 = [
@@ -2731,13 +2862,13 @@ def test_gecko_opcode_markers_not_emitted_when_disabled(self):
         ]
         collector.collect(frames2)
 
-        # opcode_state should not be tracked
-        self.assertEqual(len(collector.opcode_state), 0)
+        profile = collector._build_profile()
+        self.assertEqual(gecko_opcode_marker_data(profile), [])
+        self.assertEqual(profile["meta"]["markerSchema"], [])
 
     def test_gecko_opcode_with_none_opcode(self):
         """Test that None opcode doesn't cause issues."""
         collector = GeckoCollector(sample_interval_usec=1000, opcodes=True)
-        self.addCleanup(lambda: collector.spill_dir.cleanup())
 
         # Frame with no opcode (None)
         frame = MockFrameInfo("test.py", 10, "func", opcode=None)
@@ -2749,9 +2880,8 @@ def test_gecko_opcode_with_none_opcode(self):
         ]
         collector.collect(frames)
 
-        # Should track the state but opcode is None
-        self.assertIn(1, collector.opcode_state)
-        self.assertIsNone(collector.opcode_state[1][0])
+        profile = collector._build_profile()
+        self.assertEqual(gecko_opcode_marker_data(profile), [])
 
 
 class TestCollectorFrameFormat(unittest.TestCase):