From a1c737b73d3658be0e1d072a340d42e3d96373c6 Mon Sep 17 00:00:00 2001
From: Furkan Onder <furkanonder@protonmail.com>
Date: Mon, 31 Jul 2023 11:33:37 +0300
Subject: [PATCH] gh-104280: Add test cases for DTrace probes (#107125)

---
 Lib/test/test_dtrace.py | 82 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/Lib/test/test_dtrace.py b/Lib/test/test_dtrace.py
index 4b971deacc1..e1adf8e9748 100644
--- a/Lib/test/test_dtrace.py
+++ b/Lib/test/test_dtrace.py
@@ -3,6 +3,7 @@ import os.path
 import re
 import subprocess
 import sys
+import sysconfig
 import types
 import unittest
 
@@ -173,6 +174,87 @@ class SystemTapOptimizedTests(TraceTests, unittest.TestCase):
     backend = SystemTapBackend()
     optimize_python = 2
 
+class CheckDtraceProbes(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        if sysconfig.get_config_var('WITH_DTRACE'):
+            readelf_major_version, readelf_minor_version = cls.get_readelf_version()
+            if support.verbose:
+                print(f"readelf version: {readelf_major_version}.{readelf_minor_version}")
+        else:
+            raise unittest.SkipTest("CPython must be configured with the --with-dtrace option.")
+
+
+    @staticmethod
+    def get_readelf_version():
+        try:
+            cmd = ["readelf", "--version"]
+            proc = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                universal_newlines=True,
+            )
+            with proc:
+                version, stderr = proc.communicate()
+
+            if proc.returncode:
+                raise Exception(
+                    f"Command {' '.join(cmd)!r} failed "
+                    f"with exit code {proc.returncode}: "
+                    f"stdout={version!r} stderr={stderr!r}"
+                )
+        except OSError:
+            raise unittest.SkipTest("Couldn't find readelf on the path")
+
+        # Regex to parse:
+        # 'GNU readelf (GNU Binutils) 2.40.0\n' -> 2.40
+        match = re.search(r"^(?:GNU) readelf.*?\b(\d+)\.(\d+)", version)
+        if match is None:
+            raise unittest.SkipTest(f"Unable to parse readelf version: {version}")
+
+        return int(match.group(1)), int(match.group(2))
+
+    def get_readelf_output(self):
+        command = ["readelf", "-n", sys.executable]
+        stdout, _ = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+        ).communicate()
+        return stdout
+
+    def test_check_probes(self):
+        readelf_output = self.get_readelf_output()
+
+        available_probe_names = [
+            "Name: import__find__load__done",
+            "Name: import__find__load__start",
+            "Name: audit",
+            "Name: gc__start",
+            "Name: gc__done",
+        ]
+
+        for probe_name in available_probe_names:
+            with self.subTest(probe_name=probe_name):
+                self.assertIn(probe_name, readelf_output)
+
+    @unittest.expectedFailure
+    def test_missing_probes(self):
+        readelf_output = self.get_readelf_output()
+
+        # Missing probes will be added in the future.
+        missing_probe_names = [
+            "Name: function__entry",
+            "Name: function__return",
+            "Name: line",
+        ]
+
+        for probe_name in missing_probe_names:
+            with self.subTest(probe_name=probe_name):
+                self.assertIn(probe_name, readelf_output)
+
 
 if __name__ == '__main__':
     unittest.main()