bugfix

Author: KRRT7

codeflash/api/aiservice.py

@@ -405,7 +405,7 @@ def generate_regression_tests(  # noqa: D417
         test_timeout: int,
         trace_id: str,
         test_index: int,
-    ) -> tuple[str, str, str] | None:
+    ) -> tuple[str, str, str, str | None, str | None] | None:
         """Generate regression tests for the given function by making a request to the Django endpoint.
 
         Parameters
@@ -456,6 +456,8 @@ def generate_regression_tests(  # noqa: D417
                 response_json["generated_tests"],
                 response_json["instrumented_behavior_tests"],
                 response_json["instrumented_perf_tests"],
+                response_json.get("instrumented_behavior_source"),
+                response_json.get("instrumented_perf_source"),
             )
         try:
             error = response.json()["error"]

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -160,4 +160,3 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
         return return_value
 
     return async_wrapper
-

codeflash/code_utils/instrument_existing_tests.py

@@ -77,6 +77,10 @@ def find_and_update_line_node(
                 call_node = node
                 if isinstance(node.func, ast.Name):
                     function_name = node.func.id
+                    
+                    if self.function_object.is_async:
+                        return [test_node]
+                    
                     node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                     node.args = [
                         ast.Name(id=function_name, ctx=ast.Load()),
@@ -98,6 +102,9 @@ def find_and_update_line_node(
                 if isinstance(node.func, ast.Attribute):
                     function_to_test = node.func.attr
                     if function_to_test == self.function_object.function_name:
+                        if self.function_object.is_async:
+                            return [test_node]
+                        
                         function_name = ast.unparse(node.func)
                         node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                         node.args = [
@@ -382,7 +389,11 @@ def inject_profiling_into_existing_test(
         )
     if test_framework == "unittest":
         new_imports.append(ast.Import(names=[ast.alias(name="timeout_decorator")]))
-    tree.body = [*new_imports, create_wrapper_function(mode), *tree.body]
+    additional_functions = []
+    if not func.is_async:
+        additional_functions.append(create_wrapper_function(mode))
+    
+    tree.body = [*new_imports, *additional_functions, *tree.body]
     return True, isort.code(ast.unparse(tree), float_to_top=True)
 
 

codeflash/models/models.py

@@ -282,6 +282,8 @@ class GeneratedTests(BaseModel):
     instrumented_perf_test_source: str
     behavior_file_path: Path
     perf_file_path: Path
+    behavior_source_path: Optional[Path] = None
+    perf_source_path: Optional[Path] = None
 
 
 class GeneratedTestsList(BaseModel):
@@ -295,6 +297,8 @@ class TestFile(BaseModel):
     original_source: Optional[str] = None
     test_type: TestType
     tests_in_file: Optional[list[TestsInFile]] = None
+    behavior_source_path: Optional[Path] = None
+    perf_source_path: Optional[Path] = None
 
 
 class TestFiles(BaseModel):

codeflash/optimization/function_optimizer.py

@@ -323,6 +323,8 @@ def generate_and_instrument_tests(
                     original_source=generated_test.generated_original_test_source,
                     test_type=TestType.GENERATED_REGRESSION,
                     tests_in_file=None,  # This is currently unused. We can discover the tests in the file if needed.
+                    behavior_source_path=generated_test.behavior_source_path,
+                    perf_source_path=generated_test.perf_source_path,
                 )
             )
             logger.info(f"Generated test {i + 1}/{count_tests}:")
@@ -1022,6 +1024,8 @@ def generate_tests_and_optimizations(
                     instrumented_perf_test_source,
                     test_behavior_path,
                     test_perf_path,
+                    behavior_source_path,
+                    perf_source_path,
                 ) = res
                 tests.append(
                     GeneratedTests(
@@ -1030,6 +1034,8 @@ def generate_tests_and_optimizations(
                         instrumented_perf_test_source=instrumented_perf_test_source,
                         behavior_file_path=test_behavior_path,
                         perf_file_path=test_perf_path,
+                        behavior_source_path=behavior_source_path,
+                        perf_source_path=perf_source_path,
                     )
                 )
         if not tests:

codeflash/verification/verifier.py

@@ -27,7 +27,7 @@ def generate_tests(
     test_index: int,
     test_path: Path,
     test_perf_path: Path,
-) -> tuple[str, str, Path] | None:
+) -> tuple[str, str, str, Path, Path, Path | None, Path | None] | None:
     # TODO: Sometimes this recreates the original Class definition. This overrides and messes up the original
     #  class import. Remove the recreation of the class definition
     start_time = time.perf_counter()
@@ -43,8 +43,8 @@ def generate_tests(
         trace_id=function_trace_id,
         test_index=test_index,
     )
-    if response and isinstance(response, tuple) and len(response) == 3:
-        generated_test_source, instrumented_behavior_test_source, instrumented_perf_test_source = response
+    if response and isinstance(response, tuple) and len(response) == 5:
+        generated_test_source, instrumented_behavior_test_source, instrumented_perf_test_source, instrumented_behavior_source, instrumented_perf_source = response
         temp_run_dir = get_run_tmp_file(Path()).as_posix()
 
         instrumented_behavior_test_source = instrumented_behavior_test_source.replace(
@@ -53,6 +53,29 @@ def generate_tests(
         instrumented_perf_test_source = instrumented_perf_test_source.replace(
             "{codeflash_run_tmp_dir_client_side}", temp_run_dir
         )
+        
+        behavior_source_path = None
+        perf_source_path = None
+        
+        if instrumented_behavior_source and function_to_optimize.is_async:
+            behavior_source_path = get_run_tmp_file(Path(f"instrumented_behavior_source_{test_index}_{function_to_optimize.function_name}.py"))
+            try:
+                with behavior_source_path.open("w", encoding="utf8") as f:
+                    f.write(instrumented_behavior_source)
+                logger.debug(f"Written instrumented behavior source to {behavior_source_path}")
+            except Exception as e:
+                logger.warning(f"Failed to write instrumented behavior source file: {e}")
+                behavior_source_path = None
+                
+        if instrumented_perf_source and function_to_optimize.is_async:
+            perf_source_path = get_run_tmp_file(Path(f"instrumented_perf_source_{test_index}_{function_to_optimize.function_name}.py"))
+            try:
+                with perf_source_path.open("w", encoding="utf8") as f:
+                    f.write(instrumented_perf_source)
+                logger.debug(f"Written instrumented performance source to {perf_source_path}")
+            except Exception as e:
+                logger.warning(f"Failed to write instrumented performance source file: {e}")
+                perf_source_path = None
     else:
         logger.warning(f"Failed to generate and instrument tests for {function_to_optimize.function_name}")
         return None
@@ -64,6 +87,8 @@ def generate_tests(
         instrumented_perf_test_source,
         test_path,
         test_perf_path,
+        behavior_source_path,
+        perf_source_path,
     )