normalize for trace and replay tests too

Author: KRRT7

codeflash/benchmarking/codeflash_trace.py

@@ -4,6 +4,7 @@
 import sqlite3
 import threading
 import time
+from pathlib import Path
 from typing import Any, Callable
 
 from codeflash.picklepatch.pickle_patcher import PicklePatcher
@@ -143,12 +144,13 @@ def wrapper(*args, **kwargs) -> Any:  # noqa: ANN002, ANN003, ANN401
                 print("Pickle limit reached")
                 self._thread_local.active_functions.remove(func_id)
                 overhead_time = time.thread_time_ns() - end_time
+                normalized_file_path = Path(func.__code__.co_filename).as_posix()
                 self.function_calls_data.append(
                     (
                         func.__name__,
                         class_name,
                         func.__module__,
-                        func.__code__.co_filename,
+                        normalized_file_path,
                         benchmark_function_name,
                         benchmark_module_path,
                         benchmark_line_number,
@@ -169,12 +171,13 @@ def wrapper(*args, **kwargs) -> Any:  # noqa: ANN002, ANN003, ANN401
                 # Add to the list of function calls without pickled args. Used for timing info only
                 self._thread_local.active_functions.remove(func_id)
                 overhead_time = time.thread_time_ns() - end_time
+                normalized_file_path = Path(func.__code__.co_filename).as_posix()
                 self.function_calls_data.append(
                     (
                         func.__name__,
                         class_name,
                         func.__module__,
-                        func.__code__.co_filename,
+                        normalized_file_path,
                         benchmark_function_name,
                         benchmark_module_path,
                         benchmark_line_number,
@@ -192,12 +195,13 @@ def wrapper(*args, **kwargs) -> Any:  # noqa: ANN002, ANN003, ANN401
             # Add to the list of function calls with pickled args, to be used for replay tests
             self._thread_local.active_functions.remove(func_id)
             overhead_time = time.thread_time_ns() - end_time
+            normalized_file_path = Path(func.__code__.co_filename).as_posix()
             self.function_calls_data.append(
                 (
                     func.__name__,
                     class_name,
                     func.__module__,
-                    func.__code__.co_filename,
+                    normalized_file_path,
                     benchmark_function_name,
                     benchmark_module_path,
                     benchmark_line_number,

codeflash/benchmarking/replay_test.py

@@ -29,20 +29,25 @@ def get_next_arg_and_return(
     db = sqlite3.connect(trace_file)
     cur = db.cursor()
     limit = num_to_get
+    
+    normalized_file_path = Path(file_path).as_posix()
 
     if class_name is not None:
         cursor = cur.execute(
             "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = ? LIMIT ?",
-            (benchmark_function_name, function_name, file_path, class_name, limit),
+            (benchmark_function_name, function_name, normalized_file_path, class_name, limit),
         )
     else:
         cursor = cur.execute(
             "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = '' LIMIT ?",
-            (benchmark_function_name, function_name, file_path, limit),
+            (benchmark_function_name, function_name, normalized_file_path, limit),
         )
 
-    while (val := cursor.fetchone()) is not None:
-        yield val[9], val[10]  # pickled_args, pickled_kwargs
+    try:
+        while (val := cursor.fetchone()) is not None:
+            yield val[9], val[10]  # pickled_args, pickled_kwargs
+    finally:
+        db.close()
 
 
 def get_function_alias(module: str, function_name: str) -> str:

tests/test_pickle_patcher.py

@@ -18,6 +18,7 @@
 from codeflash.models.models import CodePosition, TestFile, TestFiles, TestingMode, TestsInFile, TestType
 from codeflash.optimization.optimizer import Optimizer
 from codeflash.verification.equivalence import compare_test_results
+import time
 
 try:
     import sqlalchemy
@@ -156,6 +157,9 @@ def test_picklepatch_with_database_connection():
     with pytest.raises(PicklePlaceholderAccessError):
         reloaded["connection"].execute("SELECT 1")
 
+    cursor.close()
+    conn.close()    
+
 
 def test_picklepatch_with_generator():
     """Test that a data structure containing a generator is replaced by
@@ -290,6 +294,7 @@ def test_run_and_parse_picklepatch() -> None:
 
         # Close the connection to allow file cleanup on Windows
         conn.close()
+        time.sleep(1)
 
         # Handle the case where function runs too fast to be measured
         unused_socket_results = function_to_results["code_to_optimize.bubble_sort_picklepatch_test_unused_socket.bubble_sort_with_unused_socket"]
@@ -326,7 +331,9 @@ def test_run_and_parse_picklepatch() -> None:
             assert actual[4] == expected[4], f"Mismatch at index {idx} for benchmark_function_name"
             assert actual[5] == expected[5], f"Mismatch at index {idx} for benchmark_module_path"
             assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number"
-        conn.close()
+            conn.close()
+            
+            time.sleep(1)
 
         # Generate replay test
         generate_replay_test(output_file, replay_tests_dir)