temp

Author: KRRT7

code_to_optimize/async_adder.py

@@ -0,0 +1,10 @@
+import asyncio
+
+
+async def async_add(a, b):
+    """Simple async function that adds two numbers."""
+    await asyncio.sleep(0.001)  # Simulate some async work
+    print(f"codeflash stdout: Adding {a} + {b}")
+    result = a + b
+    print(f"result: {result}")
+    return result

codeflash/code_utils/instrument_existing_tests.py

@@ -73,6 +73,8 @@ def find_and_update_line_node(
         self, test_node: ast.stmt, node_name: str, index: str, test_class_name: str | None = None
     ) -> Iterable[ast.stmt] | None:
         call_node = None
+        await_node = None
+        
         for node in ast.walk(test_node):
             if isinstance(node, ast.Call) and node_in_call_position(node, self.call_positions):
                 call_node = node
@@ -121,6 +123,60 @@ def find_and_update_line_node(
                         ]
                         node.keywords = call_node.keywords
                         break
+            
+            # Check for awaited function calls
+            elif isinstance(node, ast.Await) and isinstance(node.value, ast.Call) and node_in_call_position(node.value, self.call_positions):
+                call_node = node.value
+                await_node = node
+                if isinstance(call_node.func, ast.Name):
+                    function_name = call_node.func.id
+                    call_node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
+                    call_node.args = [
+                        ast.Name(id=function_name, ctx=ast.Load()),
+                        ast.Constant(value=self.module_path),
+                        ast.Constant(value=test_class_name or None),
+                        ast.Constant(value=node_name),
+                        ast.Constant(value=self.function_object.qualified_name),
+                        ast.Constant(value=index),
+                        ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
+                        *(
+                            [ast.Name(id="codeflash_cur", ctx=ast.Load()), ast.Name(id="codeflash_con", ctx=ast.Load())]
+                            if self.mode == TestingMode.BEHAVIOR
+                            else []
+                        ),
+                        *call_node.args,
+                    ]
+                    call_node.keywords = call_node.keywords
+                    # Keep the await wrapper around the modified call
+                    await_node.value = call_node
+                    break
+                if isinstance(call_node.func, ast.Attribute):
+                    function_to_test = call_node.func.attr
+                    if function_to_test == self.function_object.function_name:
+                        function_name = ast.unparse(call_node.func)
+                        call_node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
+                        call_node.args = [
+                            ast.Name(id=function_name, ctx=ast.Load()),
+                            ast.Constant(value=self.module_path),
+                            ast.Constant(value=test_class_name or None),
+                            ast.Constant(value=node_name),
+                            ast.Constant(value=self.function_object.qualified_name),
+                            ast.Constant(value=index),
+                            ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
+                            *(
+                                [
+                                    ast.Name(id="codeflash_cur", ctx=ast.Load()),
+                                    ast.Name(id="codeflash_con", ctx=ast.Load()),
+                                ]
+                                if self.mode == TestingMode.BEHAVIOR
+                                else []
+                            ),
+                            *call_node.args,
+                        ]
+                        call_node.keywords = call_node.keywords
+                        # Keep the await wrapper around the modified call
+                        await_node.value = call_node
+                        break
 
         if call_node is None:
             return None
@@ -131,9 +187,35 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
         for inner_node in ast.walk(node):
             if isinstance(inner_node, ast.FunctionDef):
                 self.visit_FunctionDef(inner_node, node.name)
+            elif isinstance(inner_node, ast.AsyncFunctionDef):
+                self.visit_AsyncFunctionDef(inner_node, node.name)
 
         return node
 
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef, test_class_name: str | None = None) -> ast.AsyncFunctionDef:
+        """Handle async function definitions by converting to sync and back."""
+        # Convert to sync FunctionDef, process it, then convert back
+        sync_node = ast.FunctionDef(
+            name=node.name,
+            args=node.args,
+            body=node.body,
+            decorator_list=node.decorator_list,
+            returns=node.returns,
+            lineno=node.lineno,
+            col_offset=node.col_offset if hasattr(node, 'col_offset') else 0
+        )
+        processed_sync = self.visit_FunctionDef(sync_node, test_class_name)
+        # Convert back to AsyncFunctionDef
+        return ast.AsyncFunctionDef(
+            name=processed_sync.name,
+            args=processed_sync.args,
+            body=processed_sync.body,
+            decorator_list=processed_sync.decorator_list,
+            returns=processed_sync.returns,
+            lineno=processed_sync.lineno,
+            col_offset=processed_sync.col_offset if hasattr(processed_sync, 'col_offset') else 0
+        )
+        
     def visit_FunctionDef(self, node: ast.FunctionDef, test_class_name: str | None = None) -> ast.FunctionDef:
         if node.name.startswith("test_"):
             did_update = False

tests/test_instrument_all_and_run.py

@@ -255,6 +255,86 @@ def test_sort():
         test_path_perf.unlink(missing_ok=True)
 
 
+def test_async_function_behavior_results() -> None:
+    """Test that async_codeflash_wrap_string is used for async functions."""
+    code = """import asyncio
+from code_to_optimize.async_adder import async_add
+
+
+async def test_async_add():
+    result = await async_add(2, 3)
+    assert result == 5"""
+
+    expected = (
+        """import gc
+import os
+import sqlite3
+import time
+
+import dill as pickle
+
+import asyncio
+from code_to_optimize.async_adder import async_add
+
+
+"""
+        + async_codeflash_wrap_string
+        + """
+async def test_async_add():
+    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
+    codeflash_iteration = os.environ['CODEFLASH_TEST_ITERATION']
+    codeflash_con = sqlite3.connect(f'{tmp_dir_path}_{{codeflash_iteration}}.sqlite')
+    codeflash_cur = codeflash_con.cursor()
+    codeflash_cur.execute('CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT)')
+    result = await codeflash_wrap(async_add, '{module_path}', None, 'test_async_add', 'async_add', '1', codeflash_loop_index, codeflash_cur, codeflash_con, 2, 3)
+    assert result == 5
+    codeflash_con.close()
+"""
+    )
+
+    test_path = (
+        Path(__file__).parent.resolve()
+        / "../code_to_optimize/tests/pytest/test_async_adder_behavior_temp.py"
+    ).resolve()
+    test_path_perf = (
+        Path(__file__).parent.resolve()
+        / "../code_to_optimize/tests/pytest/test_async_adder_perf_temp.py"
+    ).resolve()
+    fto_path = (Path(__file__).parent.resolve() / "../code_to_optimize/async_adder.py").resolve()
+    original_code = fto_path.read_text("utf-8")
+    
+    try:
+        with test_path.open("w") as f:
+            f.write(code)
+
+        tests_root = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/").resolve()
+        project_root_path = (Path(__file__).parent / "..").resolve()
+        original_cwd = Path.cwd()
+        run_cwd = Path(__file__).parent.parent.resolve()
+        func = FunctionToOptimize(function_name="async_add", parents=[], file_path=Path(fto_path), is_async=True)
+        os.chdir(run_cwd)
+        success, new_test = inject_profiling_into_existing_test(
+            test_path,
+            [CodePosition(6, 19)],
+            func,
+            project_root_path,
+            "pytest",
+            mode=TestingMode.BEHAVIOR,
+            is_async=True,
+        )
+        os.chdir(original_cwd)
+        assert success
+        assert new_test is not None
+        assert "await wrapped(*args, **kwargs)" in new_test
+        assert "async def codeflash_wrap" in new_test
+        assert "await codeflash_wrap(async_add" in new_test
+                
+    finally:
+        fto_path.write_text(original_code, "utf-8")
+        test_path.unlink(missing_ok=True)
+        test_path_perf.unlink(missing_ok=True)
+
+
 def test_class_method_full_instrumentation() -> None:
     code = """from code_to_optimize.bubble_sort_method import BubbleSorter