do instrument generated regression tests

Author: misrasaurabh1

codeflash/code_utils/instrument_existing_tests.py

@@ -632,16 +632,15 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
 
 
 def inject_async_profiling_into_existing_test(
-    test_path: Path,
+    test_string: str,
     call_positions: list[CodePosition],
     function_to_optimize: FunctionToOptimize,
     tests_project_root: Path,
     mode: TestingMode = TestingMode.BEHAVIOR,
+    test_path: Path | None = None,
 ) -> tuple[bool, str | None]:
     """Inject profiling for async function calls by setting environment variables before each call."""
-    with test_path.open(encoding="utf8") as f:
-        test_code = f.read()
-
+    test_code = test_string
     try:
         tree = ast.parse(test_code)
     except SyntaxError:
@@ -704,6 +703,7 @@ def detect_frameworks_from_code(code: str) -> dict[str, str]:
 
 
 def inject_profiling_into_existing_test(
+    test_string: str,
     test_path: Path,
     call_positions: list[CodePosition],
     function_to_optimize: FunctionToOptimize,
@@ -715,7 +715,7 @@ def inject_profiling_into_existing_test(
         from codeflash.languages.javascript.instrument import inject_profiling_into_existing_js_test
 
         return inject_profiling_into_existing_js_test(
-            test_path, call_positions, function_to_optimize, tests_project_root, mode.value
+            test_string=test_string, call_positions=call_positions, function_to_optimize=function_to_optimize, tests_project_root=tests_project_root, mode= mode.value, test_path=test_path
         )
 
     if is_java():
@@ -725,15 +725,14 @@ def inject_profiling_into_existing_test(
 
     if function_to_optimize.is_async:
         return inject_async_profiling_into_existing_test(
-            test_path, call_positions, function_to_optimize, tests_project_root, mode
+            test_string=test_string, call_positions=call_positions, function_to_optimize=function_to_optimize, tests_project_root=tests_project_root, mode=mode.value, test_path=test_path
         )
 
-    with test_path.open(encoding="utf8") as f:
-        test_code = f.read()
 
-    used_frameworks = detect_frameworks_from_code(test_code)
+
+    used_frameworks = detect_frameworks_from_code(test_string)
     try:
-        tree = ast.parse(test_code)
+        tree = ast.parse(test_string)
     except SyntaxError:
         logger.exception(f"Syntax error in code in file - {test_path}")
         return False, None

codeflash/languages/base.py

@@ -567,18 +567,20 @@ def ensure_runtime_environment(self, project_root: Path) -> bool:
 
     def instrument_existing_test(
         self,
-        test_path: Path,
+        test_string: str,
         call_positions: Sequence[Any],
         function_to_optimize: Any,
         tests_project_root: Path,
         mode: str,
+        test_path: Path | None
     ) -> tuple[bool, str | None]:
         """Inject profiling code into an existing test file.
 
         Wraps function calls with capture/benchmark instrumentation for
         behavioral verification and performance benchmarking.
 
         Args:
+            test_string: String containing the test file contents.
             test_path: Path to the test file.
             call_positions: List of code positions where the function is called.
             function_to_optimize: The function being optimized.

codeflash/languages/java/instrumentation.py

@@ -232,13 +232,11 @@ def instrument_for_benchmarking(
 
 
 def instrument_existing_test(
-    test_path: Path,
-    call_positions: Sequence,
+    test_string: str,
     function_to_optimize: Any,  # FunctionToOptimize or FunctionToOptimize
-    tests_project_root: Path,
     mode: str,  # "behavior" or "performance"
-    analyzer: JavaAnalyzer | None = None,
-    output_class_suffix: str | None = None,  # Suffix for renamed class
+    test_path: Path | None = None,
+    test_class_name: str | None = None,
 ) -> tuple[bool, str | None]:
     """Inject profiling code into an existing test file.
 
@@ -248,7 +246,7 @@ def instrument_existing_test(
     3. For performance mode: adds timing instrumentation with stdout markers
 
     Args:
-        test_path: Path to the test file.
+        test_string: String to the test file.
         call_positions: List of code positions where the function is called.
         function_to_optimize: The function being optimized.
         tests_project_root: Root directory of tests.
@@ -260,16 +258,16 @@ def instrument_existing_test(
         Tuple of (success, modified_source).
 
     """
-    try:
-        source = test_path.read_text(encoding="utf-8")
-    except Exception as e:
-        logger.exception("Failed to read test file %s: %s", test_path, e)
-        return False, f"Failed to read test file: {e}"
-
+    source = test_string
     func_name = _get_function_name(function_to_optimize)
 
     # Get the original class name from the file name
-    original_class_name = test_path.stem  # e.g., "AlgorithmsTest"
+    if test_path:
+        original_class_name = test_path.stem  # e.g., "AlgorithmsTest"
+    elif test_class_name is not None:
+        original_class_name = test_class_name
+    else:
+        raise ValueError("test_path or test_class_name must be provided")
 
     # Determine the new class name based on mode
     if mode == "behavior":
@@ -298,7 +296,7 @@ def instrument_existing_test(
         modified_source = _add_behavior_instrumentation(modified_source, original_class_name, func_name)
 
     logger.debug("Java %s testing for %s: renamed class %s -> %s", mode, func_name, original_class_name, new_class_name)
-
+    # Why return True here?
     return True, modified_source
 
 
@@ -800,6 +798,7 @@ def instrument_generated_java_test(
     function_name: str,
     qualified_name: str,
     mode: str,  # "behavior" or "performance"
+    function_to_optimize: FunctionToOptimize,
 ) -> str:
     """Instrument a generated Java test for behavior or performance testing.
 
@@ -834,26 +833,31 @@ def instrument_generated_java_test(
 
     original_class_name = class_match.group(1)
 
-    # Rename class based on mode
-    if mode == "behavior":
-        new_class_name = f"{original_class_name}__perfinstrumented"
-    else:
-        new_class_name = f"{original_class_name}__perfonlyinstrumented"
-
-    # Rename all references to the original class name in the source.
-    # This includes the class declaration, return types, constructor calls, etc.
-    modified_code = re.sub(
-        rf"\b{re.escape(original_class_name)}\b", new_class_name, test_code
-    )
 
     # For performance mode, add timing instrumentation
     # Use original class name (without suffix) in timing markers for consistency with Python
     if mode == "performance":
+
+
+        # Rename class based on mode
+        if mode == "behavior":
+            new_class_name = f"{original_class_name}__perfinstrumented"
+        else:
+            new_class_name = f"{original_class_name}__perfonlyinstrumented"
+
+        # Rename all references to the original class name in the source.
+        # This includes the class declaration, return types, constructor calls, etc.
+        modified_code = re.sub(
+            rf"\b{re.escape(original_class_name)}\b", new_class_name, test_code
+        )
+
         modified_code = _add_timing_instrumentation(
             modified_code,
             original_class_name,  # Use original name in markers, not the renamed class
             function_name,
         )
+    elif mode == "behavior":
+        _ , modified_code = instrument_existing_test(test_string=test_code, mode=mode, function_to_optimize=function_to_optimize, test_class_name=original_class_name)
 
     logger.debug("Instrumented generated Java test for %s (mode=%s)", function_name, mode)
     return modified_code

codeflash/languages/java/support.py

@@ -283,11 +283,12 @@ def ensure_runtime_environment(self, project_root: Path) -> bool:
 
     def instrument_existing_test(
         self,
-        test_path: Path,
+        test_string: str,
         call_positions: Sequence[Any],
         function_to_optimize: Any,
         tests_project_root: Path,
         mode: str,
+        test_path: Path | None
     ) -> tuple[bool, str | None]:
         """Inject profiling code into an existing test file."""
         return instrument_existing_test(

codeflash/languages/javascript/instrument.py

@@ -626,18 +626,20 @@ def transform_expect_calls(
 
 
 def inject_profiling_into_existing_js_test(
-    test_path: Path,
+    test_string: str,
     call_positions: list[CodePosition],
     function_to_optimize: FunctionToOptimize,
     tests_project_root: Path,
     mode: str = TestingMode.BEHAVIOR,
+    test_path: Path | None = None,
 ) -> tuple[bool, str | None]:
     """Inject profiling code into an existing JavaScript test file.
 
     This function wraps function calls with codeflash.capture() or codeflash.capturePerf()
     to enable behavioral verification and performance benchmarking.
 
     Args:
+        test_string: String contents of the test file.
         test_path: Path to the test file.
         call_positions: List of code positions where the function is called.
         function_to_optimize: The function being optimized.
@@ -648,13 +650,7 @@ def inject_profiling_into_existing_js_test(
         Tuple of (success, instrumented_code).
 
     """
-    try:
-        with test_path.open(encoding="utf8") as f:
-            test_code = f.read()
-    except Exception as e:
-        logger.error(f"Failed to read test file {test_path}: {e}")
-        return False, None
-
+    test_code = test_string
     # Get the relative path for test identification
     try:
         rel_path = test_path.relative_to(tests_project_root)

codeflash/languages/javascript/support.py

@@ -1936,11 +1936,12 @@ def ensure_runtime_environment(self, project_root: Path) -> bool:
 
     def instrument_existing_test(
         self,
-        test_path: Path,
+        test_string: str,
         call_positions: Sequence[Any],
         function_to_optimize: Any,
         tests_project_root: Path,
         mode: str,
+        test_path: Path|None,
     ) -> tuple[bool, str | None]:
         """Inject profiling code into an existing JavaScript test file.
 
@@ -1961,6 +1962,7 @@ def instrument_existing_test(
         from codeflash.languages.javascript.instrument import inject_profiling_into_existing_js_test
 
         return inject_profiling_into_existing_js_test(
+            test_string=test_string,
             test_path=test_path,
             call_positions=list(call_positions),
             function_to_optimize=function_to_optimize,

codeflash/optimization/function_optimizer.py

@@ -1849,13 +1849,15 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
                     logger.debug(f"Failed to instrument test file {test_file} for behavior testing")
                     continue
 
-                success, injected_perf_test = self.language_support.instrument_existing_test(
-                    test_path=path_obj_test_file,
-                    call_positions=[test.position for test in tests_in_file_list],
-                    function_to_optimize=self.function_to_optimize,
-                    tests_project_root=self.test_cfg.tests_project_rootdir,
-                    mode="performance",
-                )
+                with path_obj_test_file.open("r", encoding="utf8") as f:
+                    injected_behavior_test_source = f.read()
+                    success, injected_perf_test = self.language_support.instrument_existing_test(
+                        test_string=injected_behavior_test_source,
+                        call_positions=[test.position for test in tests_in_file_list],
+                        function_to_optimize=self.function_to_optimize,
+                        tests_project_root=self.test_cfg.tests_project_rootdir,
+                        mode="performance",
+                    )
                 if not success:
                     logger.debug(f"Failed to instrument test file {test_file} for performance testing")
                     continue

codeflash/verification/verifier.py

@@ -109,7 +109,7 @@ def generate_tests(
 
             # Instrument for behavior verification (renames class)
             instrumented_behavior_test_source = instrument_generated_java_test(
-                test_code=generated_test_source, function_name=func_name, qualified_name=qualified_name, mode="behavior"
+                test_code=generated_test_source, function_name=func_name, qualified_name=qualified_name, mode="behavior", function_to_optimize=function_to_optimize
             )
 
             # Instrument for performance measurement (adds timing markers)
@@ -118,6 +118,7 @@ def generate_tests(
                 function_name=func_name,
                 qualified_name=qualified_name,
                 mode="performance",
+                function_to_optimize=function_to_optimize
             )
 
             logger.debug(f"Instrumented Java tests locally for {func_name}")