@@ -33,11 +33,17 @@ def normalize_trace_output(output):
3333 result = [
3434 row .split ("\t " )
3535 for row in output .splitlines ()
36- if row and not row .startswith ('#' )
36+ if row and not row .startswith ('#' ) and not row . startswith ( '@' )
3737 ]
3838 result .sort (key = lambda row : int (row [0 ]))
3939 result = [row [1 ] for row in result ]
40- return "\n " .join (result )
40+ # Normalize paths to basenames (bpftrace outputs full paths)
41+ normalized = []
42+ for line in result :
43+ # Replace full paths with just the filename
44+ line = re .sub (r'/[^:]+/([^/:]+\.py)' , r'\1' , line )
45+ normalized .append (line )
46+ return "\n " .join (normalized )
4147 except (IndexError , ValueError ):
4248 raise AssertionError (
4349 "tracer produced unparsable output:\n {}" .format (output )
@@ -96,13 +102,186 @@ def assert_usable(self):
96102class DTraceBackend (TraceBackend ):
97103 EXTENSION = ".d"
98104 COMMAND = ["dtrace" , "-q" , "-s" ]
105+ if sys .platform == "sunos5" :
106+ COMMAND .insert (2 , "-Z" )
99107
100108
101109class SystemTapBackend (TraceBackend ):
102110 EXTENSION = ".stp"
103111 COMMAND = ["stap" , "-g" ]
104112
105113
114+ class BPFTraceBackend (TraceBackend ):
115+ EXTENSION = ".bt"
116+ COMMAND = ["bpftrace" ]
117+
118+ # Inline bpftrace programs for each test case
119+ PROGRAMS = {
120+ "call_stack" : """
121+ usdt:{python}:python:function__entry {{
122+ printf("%lld\\ tfunction__entry:%s:%s:%d\\ n",
123+ nsecs, str(arg0), str(arg1), arg2);
124+ }}
125+ usdt:{python}:python:function__return {{
126+ printf("%lld\\ tfunction__return:%s:%s:%d\\ n",
127+ nsecs, str(arg0), str(arg1), arg2);
128+ }}
129+ """ ,
130+ "gc" : """
131+ usdt:{python}:python:function__entry {{
132+ if (str(arg1) == "start") {{ @tracing = 1; }}
133+ }}
134+ usdt:{python}:python:function__return {{
135+ if (str(arg1) == "start") {{ @tracing = 0; }}
136+ }}
137+ usdt:{python}:python:gc__start {{
138+ if (@tracing) {{
139+ printf("%lld\\ tgc__start:%d\\ n", nsecs, arg0);
140+ }}
141+ }}
142+ usdt:{python}:python:gc__done {{
143+ if (@tracing) {{
144+ printf("%lld\\ tgc__done:%lld\\ n", nsecs, arg0);
145+ }}
146+ }}
147+ END {{ clear(@tracing); }}
148+ """ ,
149+ }
150+
151+ # Which test scripts to filter by filename (None = use @tracing flag)
152+ FILTER_BY_FILENAME = {"call_stack" : "call_stack.py" }
153+
154+ @staticmethod
155+ def _filter_probe_rows (output ):
156+ return "\n " .join (
157+ line for line in output .splitlines ()
158+ if line .partition ("\t " )[0 ].isdigit ()
159+ )
160+
161+ # Expected outputs for each test case
162+ # Note: bpftrace captures <module> entry/return and may have slight timing
163+ # differences compared to SystemTap due to probe firing order
164+ EXPECTED = {
165+ "call_stack" : """function__entry:call_stack.py:<module>:0
166+ function__entry:call_stack.py:start:23
167+ function__entry:call_stack.py:function_1:1
168+ function__entry:call_stack.py:function_3:9
169+ function__return:call_stack.py:function_3:10
170+ function__return:call_stack.py:function_1:2
171+ function__entry:call_stack.py:function_2:5
172+ function__entry:call_stack.py:function_1:1
173+ function__entry:call_stack.py:function_3:9
174+ function__return:call_stack.py:function_3:10
175+ function__return:call_stack.py:function_1:2
176+ function__return:call_stack.py:function_2:6
177+ function__entry:call_stack.py:function_3:9
178+ function__return:call_stack.py:function_3:10
179+ function__entry:call_stack.py:function_4:13
180+ function__return:call_stack.py:function_4:14
181+ function__entry:call_stack.py:function_5:18
182+ function__return:call_stack.py:function_5:21
183+ function__return:call_stack.py:start:28
184+ function__return:call_stack.py:<module>:30""" ,
185+ "gc" : """gc__start:0
186+ gc__done:0
187+ gc__start:1
188+ gc__done:0
189+ gc__start:2
190+ gc__done:0
191+ gc__start:2
192+ gc__done:1""" ,
193+ }
194+
195+ def run_case (self , name , optimize_python = None ):
196+ if name not in self .PROGRAMS :
197+ raise unittest .SkipTest (f"No bpftrace program for { name } " )
198+
199+ python_file = abspath (name + ".py" )
200+ python_flags = []
201+ if optimize_python :
202+ python_flags .extend (["-O" ] * optimize_python )
203+
204+ subcommand = [sys .executable ] + python_flags + [python_file ]
205+ program = self .PROGRAMS [name ].format (python = sys .executable )
206+
207+ try :
208+ proc = subprocess .Popen (
209+ ["bpftrace" , "-e" , program , "-c" , " " .join (subcommand )],
210+ stdout = subprocess .PIPE ,
211+ stderr = subprocess .PIPE ,
212+ universal_newlines = True ,
213+ )
214+ stdout , stderr = proc .communicate (timeout = 60 )
215+ except subprocess .TimeoutExpired :
216+ proc .kill ()
217+ raise AssertionError ("bpftrace timed out" )
218+ except (FileNotFoundError , PermissionError ) as e :
219+ raise unittest .SkipTest (f"bpftrace not available: { e } " )
220+
221+ if proc .returncode != 0 :
222+ raise AssertionError (
223+ f"bpftrace failed with code { proc .returncode } :\n { stderr } "
224+ )
225+
226+ stdout = self ._filter_probe_rows (stdout )
227+
228+ # Filter output by filename if specified (bpftrace captures everything)
229+ if name in self .FILTER_BY_FILENAME :
230+ filter_filename = self .FILTER_BY_FILENAME [name ]
231+ filtered_lines = [
232+ line for line in stdout .splitlines ()
233+ if filter_filename in line
234+ ]
235+ stdout = "\n " .join (filtered_lines )
236+
237+ actual_output = normalize_trace_output (stdout )
238+ expected_output = self .EXPECTED [name ].strip ()
239+
240+ return (expected_output , actual_output )
241+
242+ def assert_usable (self ):
243+ # Check if bpftrace is available and can attach to USDT probes
244+ program = f'usdt:{ sys .executable } :python:function__entry {{ printf("probe: success\\ n"); exit(); }}'
245+ try :
246+ proc = subprocess .Popen (
247+ ["bpftrace" , "-e" , program , "-c" , f"{ sys .executable } -c pass" ],
248+ stdout = subprocess .PIPE ,
249+ stderr = subprocess .PIPE ,
250+ universal_newlines = True ,
251+ )
252+ stdout , stderr = proc .communicate (timeout = 10 )
253+ except subprocess .TimeoutExpired :
254+ proc .kill ()
255+ proc .communicate () # Clean up
256+ raise unittest .SkipTest ("bpftrace timed out during usability check" )
257+ except OSError as e :
258+ raise unittest .SkipTest (f"bpftrace not available: { e } " )
259+
260+ # Check for permission errors (bpftrace usually requires root)
261+ if proc .returncode != 0 :
262+ raise unittest .SkipTest (
263+ f"bpftrace(1) failed with code { proc .returncode } : { stderr } "
264+ )
265+
266+ if "probe: success" not in stdout :
267+ raise unittest .SkipTest (
268+ f"bpftrace(1) failed: stdout={ stdout !r} stderr={ stderr !r} "
269+ )
270+
271+
272+ class BPFTraceOutputTests (unittest .TestCase ):
273+ def test_filter_probe_rows_ignores_warnings (self ):
274+ output = """stdin:1-19: WARNING: found external warnings
275+ HINT: include/vmlinux.h:1439:3: warning: declaration does not declare anything
276+ 4623214882928\t gc__start:0
277+ 4623214885575\t gc__done:0
278+ """
279+ self .assertEqual (
280+ BPFTraceBackend ._filter_probe_rows (output ),
281+ "4623214882928\t gc__start:0\n 4623214885575\t gc__done:0" ,
282+ )
283+
284+
106285@unittest .skipIf (MS_WINDOWS , "Tests not compliant with trace on Windows." )
107286class TraceTests :
108287 # unittest.TestCase options
@@ -127,7 +306,8 @@ def test_function_entry_return(self):
127306 def test_verify_call_opcodes (self ):
128307 """Ensure our call stack test hits all function call opcodes"""
129308
130- opcodes = set (["CALL_FUNCTION" , "CALL_FUNCTION_EX" , "CALL_FUNCTION_KW" ])
309+ # Modern Python uses CALL, CALL_KW, and CALL_FUNCTION_EX
310+ opcodes = set (["CALL" , "CALL_FUNCTION_EX" , "CALL_KW" ])
131311
132312 with open (abspath ("call_stack.py" )) as f :
133313 code_string = f .read ()
@@ -152,9 +332,6 @@ def get_function_instructions(funcname):
152332 def test_gc (self ):
153333 self .run_case ("gc" )
154334
155- def test_line (self ):
156- self .run_case ("line" )
157-
158335
159336class DTraceNormalTests (TraceTests , unittest .TestCase ):
160337 backend = DTraceBackend ()
@@ -175,6 +352,17 @@ class SystemTapOptimizedTests(TraceTests, unittest.TestCase):
175352 backend = SystemTapBackend ()
176353 optimize_python = 2
177354
355+
356+ class BPFTraceNormalTests (TraceTests , unittest .TestCase ):
357+ backend = BPFTraceBackend ()
358+ optimize_python = 0
359+
360+
361+ class BPFTraceOptimizedTests (TraceTests , unittest .TestCase ):
362+ backend = BPFTraceBackend ()
363+ optimize_python = 2
364+
365+
178366class CheckDtraceProbes (unittest .TestCase ):
179367 @classmethod
180368 def setUpClass (cls ):
@@ -235,6 +423,8 @@ def test_check_probes(self):
235423 "Name: audit" ,
236424 "Name: gc__start" ,
237425 "Name: gc__done" ,
426+ "Name: function__entry" ,
427+ "Name: function__return" ,
238428 ]
239429
240430 for probe_name in available_probe_names :
@@ -247,8 +437,6 @@ def test_missing_probes(self):
247437
248438 # Missing probes will be added in the future.
249439 missing_probe_names = [
250- "Name: function__entry" ,
251- "Name: function__return" ,
252440 "Name: line" ,
253441 ]
254442
0 commit comments