test

Author: KRRT7

codeflash/__init__.py

@@ -0,0 +1,5 @@
+# import pytest
+
+# pytest.register_assert_rewrite("codeflash.benchmarking")
+# pytest.register_assert_rewrite("codeflash.tracing")
+# pytest.register_assert_rewrite("codeflash.code_utils")

codeflash/benchmarking/pytest_new_process_trace_benchmarks.py

@@ -1,17 +1,16 @@
 import sys
 from pathlib import Path
 
+import pytest
+pytest.register_assert_rewrite("codeflash")
+
 benchmarks_root = sys.argv[1]
 tests_root = sys.argv[2]
 trace_file = sys.argv[3]
 # current working directory
 project_root = Path.cwd()
 
 if __name__ == "__main__":
-    import pytest
-
-    pytest.register_assert_rewrite("codeflash")
-
     from codeflash.benchmarking.codeflash_trace import codeflash_trace
     from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin