Fix tests

NVIDIA · Nov 12, 2020 · 064462b · 064462b
1 parent 0187a22
commit 064462b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/pyprof/nvtx/dlprof.py b/pyprof/nvtx/dlprof.py
@@ -45,8 +45,9 @@ def should_skip_frame_name(cls, name, prev_name):
         #    Remove back-to-back duplicates of the same function name.
         #    This is common during recursion
         #
-        if name in ["wrapper_func", "always_benchmark_wrapper"]:
-            return True
+        for prefix in ["wrapper_func", "always_benchmark_wrapper"]:
+            if name.startswith(prefix):
+                return True
         if name == prev_name:
             return True
         return False

diff --git a/qa/L0_function_stack/test_pyprof_func_stack.py b/qa/L0_function_stack/test_pyprof_func_stack.py
@@ -114,7 +114,7 @@ def func1():
 
         func1()
 
-    # Test that lambdas are ignored in hierarchy
+    # Test that lambdas are NOT ignored in hierarchy
     # Function stack is func1->func2->lambda->func3->verify
     # Local function 'verify' gets recognized as a member of TestPyProfFuncStack because it uses 'self'
     #
@@ -124,7 +124,7 @@ def verify():
             tracemarker = pyprof.nvtx.nvmarker.traceMarker("opname")
             self.compare_funcstack(
                 tracemarker,
-                "TestPyProfFuncStack::test_ignore_lambda/func1/func2/func3/TestPyProfFuncStack::verify/opname"
+                "TestPyProfFuncStack::test_ignore_lambda/func1/func2/<lambda>/func3/TestPyProfFuncStack::verify/opname"
             )
 
         def func3():