Fix test_perf_maps_generation

kenodegard · kenodegard · commit 2b026dd7d52f · 2024-02-14T08:21:04.000-06:00
diff --git a/src/pytest_codspeed/plugin.py b/src/pytest_codspeed/plugin.py
@@ -172,23 +172,26 @@ def wrap_pyfunc_with_instrumentation(
     lib: LibType,
     nodeid: str,
     config: pytest.Config,
-    fn: Callable[P, T],
+    testfunction: Callable[P, T],
 ) -> Callable[P, T]:
-    @functools.wraps(fn)
+    @functools.wraps(testfunction)
     def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        def __codspeed_root_frame__():
+            return testfunction(*args, **kwargs)
+
         is_gc_enabled = gc.isenabled()
         if is_gc_enabled:
             gc.collect()
             gc.disable()
         try:
             if SUPPORTS_PERF_TRAMPOLINE:
                 # Warmup CPython performance map cache
-                fn(*args, **kwargs)
+                __codspeed_root_frame__()
 
             lib.zero_stats()
             lib.start_instrumentation()
             try:
-                return fn(*args, **kwargs)
+                return __codspeed_root_frame__()
             finally:
                 lib.stop_instrumentation()
                 uri = get_git_relative_uri(nodeid, config.rootpath)
diff --git a/tests/test_pytest_plugin.py b/tests/test_pytest_plugin.py
@@ -306,8 +306,10 @@ def fixtured_child():
 
     with open(perf_filepath) as perf_file:
         lines = perf_file.readlines()
+
         assert any(
-            "py::_run_with_instrumentation.<locals>.__codspeed_root_frame__" in line
+            "py::wrap_pyfunc_with_instrumentation.<locals>.wrapper.<locals>.__codspeed_root_frame__"
+            in line
             for line in lines
         ), "No root frame found in perf map"
         assert any(
