async tests instrumentation

Author: KRRT7

code_to_optimize/code_directories/async_repo/pyproject.toml

@@ -0,0 +1,7 @@
+[tool.codeflash]
+# All paths are relative to this pyproject.toml's directory.
+module-root = "."
+tests-root = "tests"
+test-framework = "pytest"
+ignore-paths = []
+formatter-cmds = ["black $file"]

code_to_optimize/code_directories/async_repo/shocker.py

@@ -0,0 +1,5 @@
+from time import sleep
+
+async def tasked():
+    sleep(0.002)
+    return "Tasked"

code_to_optimize/code_directories/async_repo/tests/test_shocker.py

@@ -0,0 +1,26 @@
+import pytest
+import asyncio
+from shocker import tasked
+
+
+@pytest.mark.asyncio
+async def test_tasked_basic():
+    result = await tasked()
+    assert result == "Tasked"
+
+
+@pytest.mark.asyncio
+async def test_tasked_gather():
+    results = await asyncio.gather(*(tasked() for _ in range(5)))
+    assert results == ["Tasked"] * 5
+
+
+def test_tasked_many_parallel_invocations():
+    async def run_many():
+        tasks = [tasked() for _ in range(1000)]
+        results = await asyncio.gather(*tasks)
+        return results
+
+    results = asyncio.run(run_many())
+    assert len(results) == 1000, "Should return 1000 results"
+    assert all(r == "Tasked" for r in results), "All results should be 'Tasked'"

codeflash/code_utils/instrument_existing_tests.py

@@ -48,6 +48,39 @@ def is_argument_name(name: str, arguments_node: ast.arguments) -> bool:
     )
 
 
+class AsyncIOGatherRemover(ast.NodeTransformer):
+    def _contains_asyncio_gather(self, node: ast.AST) -> bool:
+        """Check if a node contains asyncio.gather calls."""
+        for child_node in ast.walk(node):
+            if (
+                isinstance(child_node, ast.Call)
+                and isinstance(child_node.func, ast.Attribute)
+                and isinstance(child_node.func.value, ast.Name)
+                and child_node.func.value.id == "asyncio"
+                and child_node.func.attr == "gather"
+            ):
+                return True
+
+            if (
+                isinstance(child_node, ast.Call)
+                and isinstance(child_node.func, ast.Name)
+                and child_node.func.id == "gather"
+            ):
+                return True
+
+        return False
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef | None:
+        if node.name.startswith("test_") and self._contains_asyncio_gather(node):
+            return None
+        return self.generic_visit(node)
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef | None:
+        if node.name.startswith("test_") and self._contains_asyncio_gather(node):
+            return None
+        return self.generic_visit(node)
+
+
 class InjectPerfOnly(ast.NodeTransformer):
     def __init__(
         self,
@@ -397,6 +430,7 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                         file_path=self.function.file_path,
                         starting_line=self.function.starting_line,
                         ending_line=self.function.ending_line,
+                        is_async=self.function.is_async,
                     )
                 else:
                     self.imported_as = FunctionToOptimize(
@@ -405,6 +439,7 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                         file_path=self.function.file_path,
                         starting_line=self.function.starting_line,
                         ending_line=self.function.ending_line,
+                        is_async=self.function.is_async,
                     )
 
 
@@ -415,7 +450,6 @@ def inject_profiling_into_existing_test(
     tests_project_root: Path,
     test_framework: str,
     mode: TestingMode = TestingMode.BEHAVIOR,
-    is_async: bool = False,
 ) -> tuple[bool, str | None]:
     with test_path.open(encoding="utf8") as f:
         test_code = f.read()
@@ -430,6 +464,13 @@ def inject_profiling_into_existing_test(
     import_visitor.visit(tree)
     func = import_visitor.imported_as
 
+    is_async = function_to_optimize.is_async
+    logger.debug(f"Using async status from discovery phase for {function_to_optimize.function_name}: {is_async}")
+
+    if is_async:
+        asyncio_gather_remover = AsyncIOGatherRemover()
+        tree = asyncio_gather_remover.visit(tree)
+
     tree = InjectPerfOnly(func, test_module_path, test_framework, call_positions, mode=mode, is_async=is_async).visit(
         tree
     )
@@ -444,11 +485,15 @@ def inject_profiling_into_existing_test(
         )
     if test_framework == "unittest":
         new_imports.append(ast.Import(names=[ast.alias(name="timeout_decorator")]))
+    if is_async:
+        new_imports.append(ast.Import(names=[ast.alias(name="inspect")]))
     tree.body = [*new_imports, create_wrapper_function(mode, is_async), *tree.body]
     return True, isort.code(ast.unparse(tree), float_to_top=True)
 
 
-def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR, is_async: bool = False) -> ast.FunctionDef:
+def create_wrapper_function(
+    mode: TestingMode = TestingMode.BEHAVIOR, is_async: bool = False
+) -> ast.FunctionDef | ast.AsyncFunctionDef:
     lineno = 1
     wrapper_body: list[ast.stmt] = [
         ast.Assign(
@@ -624,22 +669,70 @@ def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR, is_async:
                     ),
                     lineno=lineno + 11,
                 ),
-                ast.Assign(
-                    targets=[ast.Name(id="return_value", ctx=ast.Store())],
-                    value=ast.Await(
-                        value=ast.Call(
-                            func=ast.Name(id="wrapped", ctx=ast.Load()),
-                            args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
-                            keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
-                        )
-                    )
+                # For async wrappers
+                # Call the wrapped function first, then check if result is awaitable before awaiting.
+                # This handles mixed scenarios where async tests might call both sync and async functions.
+                *(
+                    [
+                        ast.Assign(
+                            targets=[ast.Name(id="ret", ctx=ast.Store())],
+                            value=ast.Call(
+                                func=ast.Name(id="wrapped", ctx=ast.Load()),
+                                args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
+                                keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
+                            ),
+                            lineno=lineno + 12,
+                        ),
+                        ast.If(
+                            test=ast.Call(
+                                func=ast.Attribute(
+                                    value=ast.Name(id="inspect", ctx=ast.Load()), attr="isawaitable", ctx=ast.Load()
+                                ),
+                                args=[ast.Name(id="ret", ctx=ast.Load())],
+                                keywords=[],
+                            ),
+                            body=[
+                                ast.Assign(
+                                    targets=[ast.Name(id="counter", ctx=ast.Store())],
+                                    value=ast.Call(
+                                        func=ast.Attribute(
+                                            value=ast.Name(id="time", ctx=ast.Load()),
+                                            attr="perf_counter_ns",
+                                            ctx=ast.Load(),
+                                        ),
+                                        args=[],
+                                        keywords=[],
+                                    ),
+                                    lineno=lineno + 14,
+                                ),
+                                ast.Assign(
+                                    targets=[ast.Name(id="return_value", ctx=ast.Store())],
+                                    value=ast.Await(value=ast.Name(id="ret", ctx=ast.Load())),
+                                    lineno=lineno + 15,
+                                ),
+                            ],
+                            orelse=[
+                                ast.Assign(
+                                    targets=[ast.Name(id="return_value", ctx=ast.Store())],
+                                    value=ast.Name(id="ret", ctx=ast.Load()),
+                                    lineno=lineno + 16,
+                                )
+                            ],
+                            lineno=lineno + 13,
+                        ),
+                    ]
                     if is_async
-                    else ast.Call(
-                        func=ast.Name(id="wrapped", ctx=ast.Load()),
-                        args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
-                        keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
-                    ),
-                    lineno=lineno + 12,
+                    else [
+                        ast.Assign(
+                            targets=[ast.Name(id="return_value", ctx=ast.Store())],
+                            value=ast.Call(
+                                func=ast.Name(id="wrapped", ctx=ast.Load()),
+                                args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
+                                keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
+                            ),
+                            lineno=lineno + 12,
+                        )
+                    ]
                 ),
                 ast.Assign(
                     targets=[ast.Name(id="codeflash_duration", ctx=ast.Store())],

codeflash/discovery/functions_to_optimize.py

@@ -75,7 +75,7 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
             parents: CSTNode | None = self.get_metadata(cst.metadata.ParentNodeProvider, node)
             ast_parents: list[FunctionParent] = []
             while parents is not None:
-                if isinstance(parents, (cst.FunctionDef, cst.ClassDef)):
+                if isinstance(parents, (cst.FunctionDef, cst.AsyncFunctionDef, cst.ClassDef)):
                     ast_parents.append(FunctionParent(parents.name.value, parents.__class__.__name__))
                 parents = self.get_metadata(cst.metadata.ParentNodeProvider, parents, default=None)
             self.functions.append(
@@ -85,6 +85,29 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
                     parents=list(reversed(ast_parents)),
                     starting_line=pos.start.line,
                     ending_line=pos.end.line,
+                    is_async=False,
+                )
+            )
+
+    def visit_AsyncFunctionDef(self, node: cst.AsyncFunctionDef) -> None:
+        return_visitor: ReturnStatementVisitor = ReturnStatementVisitor()
+        node.visit(return_visitor)
+        if return_visitor.has_return_statement:
+            pos: CodeRange = self.get_metadata(cst.metadata.PositionProvider, node)
+            parents: CSTNode | None = self.get_metadata(cst.metadata.ParentNodeProvider, node)
+            ast_parents: list[FunctionParent] = []
+            while parents is not None:
+                if isinstance(parents, (cst.FunctionDef, cst.AsyncFunctionDef, cst.ClassDef)):
+                    ast_parents.append(FunctionParent(parents.name.value, parents.__class__.__name__))
+                parents = self.get_metadata(cst.metadata.ParentNodeProvider, parents, default=None)
+            self.functions.append(
+                FunctionToOptimize(
+                    function_name=node.name.value,
+                    file_path=self.file_path,
+                    parents=list(reversed(ast_parents)),
+                    starting_line=pos.start.line,
+                    ending_line=pos.end.line,
+                    is_async=True,
                 )
             )
 

tests/test_instrument_all_and_run.py

@@ -62,7 +62,12 @@
     gc.disable()
     try:
         counter = time.perf_counter_ns()
-        return_value = await wrapped(*args, **kwargs)
+        ret = wrapped(*args, **kwargs)
+        if inspect.isawaitable(ret):
+            counter = time.perf_counter_ns()
+            return_value = await ret
+        else:
+            return_value = ret
         codeflash_duration = time.perf_counter_ns() - counter
     except Exception as e:
         codeflash_duration = time.perf_counter_ns() - counter
@@ -266,14 +271,15 @@ async def test_async_add():
     assert result == 5"""
 
     expected = (
-        """import gc
+        """import asyncio
+import gc
+import inspect
 import os
 import sqlite3
 import time
 
 import dill as pickle
 
-import asyncio
 from code_to_optimize.async_adder import async_add
 
 
@@ -286,7 +292,7 @@ async def test_async_add():
     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)
+    result = await codeflash_wrap(async_add, '{module_path}', None, 'test_async_add', 'async_add', '0', codeflash_loop_index, codeflash_cur, codeflash_con, 2, 3)
     assert result == 5
     codeflash_con.close()
 """
@@ -320,15 +326,14 @@ async def test_async_add():
             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
-                
+        assert new_test.replace('"', "'") == expected.format(
+            module_path="code_to_optimize.tests.pytest.test_async_adder_behavior_temp", tmp_dir_path=get_run_tmp_file(Path("test_return_values"))
+        ).replace('"', "'")
     finally:
         fto_path.write_text(original_code, "utf-8")
         test_path.unlink(missing_ok=True)