test testbench

Author: KRRT7

tests/scripts/end_to_end_test_tracer_replay.py

@@ -10,7 +10,7 @@ def run_test(expected_improvement_pct: int) -> bool:
         min_improvement_x=0.1,
         expected_unit_tests=1,
         coverage_expectations=[
-            CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[2, 3, 4, 6, 9]),
+            # CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[2, 3, 4, 6, 9]),
             CoverageExpectation(function_name="add_numbers", expected_coverage=100.0, expected_lines=[2, 3, 4, 6, 9]),
         ],
     )

tests/scripts/end_to_end_test_utilities.py

@@ -80,7 +80,8 @@ def run_codeflash_command(
 ) -> bool:
     logging.basicConfig(level=logging.INFO)
     if config.trace_mode:
-        return run_trace_test(cwd, config, expected_improvement_pct)
+        # return run_trace_test(cwd, config, expected_improvement_pct)
+        return run_trace_test2(cwd, config, expected_improvement_pct)
 
     path_to_file = cwd / config.file_path
     file_contents = path_to_file.read_text("utf-8")
@@ -228,6 +229,55 @@ def run_trace_test(cwd: pathlib.Path, config: TestConfig, expected_improvement_p
     return validate_output(stdout, return_code, expected_improvement_pct, config)
 
 
+def run_trace_test2(cwd: pathlib.Path, config: TestConfig, expected_improvement_pct: int) -> bool:
+    # First command: Run the tracer
+    test_root = cwd / "tests" / (config.test_framework or "")
+    clear_directory(test_root)
+    command = ["python", "-m", "codeflash.tracer", "-o", "codeflash.trace", "testbench.py"]
+    process = subprocess.Popen(
+        command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=str(cwd), env=os.environ.copy()
+    )
+
+    output = []
+    for line in process.stdout:
+        logging.info(line.strip())
+        output.append(line)
+
+    return_code = process.wait()
+    stdout = "".join(output)
+
+    if return_code != 0:
+        logging.error(f"Tracer command returned exit code {return_code}")
+        return False
+
+    functions_traced = re.search(r"Traced (\d+) function calls successfully and replay test created at - (.*)$", stdout)
+    if not functions_traced or int(functions_traced.group(1)) != 3:
+        logging.error("Expected 3 traced functions")
+        return False
+
+    replay_test_path = pathlib.Path(functions_traced.group(2))
+    if not replay_test_path.exists():
+        logging.error(f"Replay test file missing at {replay_test_path}")
+        return False
+
+    # Second command: Run optimization
+    command = ["python", "../../../codeflash/main.py", "--replay-test", str(replay_test_path), "--no-pr"]
+    process = subprocess.Popen(
+        command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=str(cwd), env=os.environ.copy()
+    )
+
+    output = []
+    for line in process.stdout:
+        logging.info(line.strip())
+        output.append(line)
+
+    return_code = process.wait()
+    stdout = "".join(output)
+
+    return validate_output(stdout, return_code, expected_improvement_pct, config)
+
+
+
 def run_with_retries(test_func, *args, **kwargs) -> bool:
     max_retries = int(os.getenv("MAX_RETRIES", 3))
     retry_delay = int(os.getenv("RETRY_DELAY", 5))