fix test

NVIDIA · Jul 27, 2020 · e7301d7 · e7301d7
1 parent 8b91cb0
commit e7301d7
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/qa/L0_function_stack/test_pyprof_func_stack.py b/qa/L0_function_stack/test_pyprof_func_stack.py
@@ -37,16 +37,19 @@ def tearDown(self):
         pass
 
     def compare_funcstack(self, actual_tracemarker, expected_str):
-        # Given a funcstack string, remove TestPyProfFuncStack::run and everything above it
+        # Given a funcstack string, remove TestPyProfFuncStack::__call__/run and everything above it
         #
         def remove_test_class_hierarchy(x):
             separator = "/"
             fn_split = x.split(separator)
+            split = 0
+            # Find the LAST instance of run in the split
+            #
             for i, n in enumerate(fn_split):
                 if (n == "run"):
-                    fn_split = fn_split[i+1:]
-                    break
+                    split = i+1
 
+            fn_split = fn_split[split:]
             joined = separator.join(fn_split)
             return joined
 
@@ -103,7 +106,7 @@ def test_ignore_class_call(self):
 
         def wrapper_func():
             tracemarker = pyprof.nvtx.nvmarker.traceMarker("opname")
-            self.compare_funcstack(tracemarker,"test_ignore_class_call/func1/func2/func3/opname")
+            self.compare_funcstack(tracemarker,"test_ignore_class_call/func1/func2/ExecutableClass::__call__/func3/opname")
 
         def func3():
             wrapper_func()