Merge branch 'main' into formatter-output-fix

Author: aseembits93

codeflash/code_utils/config_consts.py

@@ -8,3 +8,4 @@
 N_TESTS_TO_GENERATE = 2
 TOTAL_LOOPING_TIME = 10.0  # 10 second candidate benchmarking budget
 COVERAGE_THRESHOLD = 60.0
+MIN_TESTCASE_PASSED_THRESHOLD = 6

codeflash/discovery/functions_to_optimize.py

@@ -94,8 +94,6 @@ def visit_FunctionDef(self, node: FunctionDef) -> None:
             self.functions.append(
                 FunctionToOptimize(function_name=node.name, file_path=self.file_path, parents=self.ast_path[:])
             )
-        # Continue visiting the body of the function to find nested functions
-        self.generic_visit(node)
 
     def generic_visit(self, node: ast.AST) -> None:
         if isinstance(node, (FunctionDef, AsyncFunctionDef, ClassDef)):

codeflash/result/critic.py

@@ -4,7 +4,11 @@
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
-from codeflash.code_utils.config_consts import COVERAGE_THRESHOLD, MIN_IMPROVEMENT_THRESHOLD
+from codeflash.code_utils.config_consts import (
+    COVERAGE_THRESHOLD,
+    MIN_IMPROVEMENT_THRESHOLD,
+    MIN_TESTCASE_PASSED_THRESHOLD,
+)
 from codeflash.models.models import TestType
 
 if TYPE_CHECKING:
@@ -50,7 +54,7 @@ def quantity_of_tests_critic(candidate_result: OptimizedCandidateResult) -> bool
     for test_type in report:
         pass_count += report[test_type]["passed"]
 
-    if pass_count >= 4:
+    if pass_count >= MIN_TESTCASE_PASSED_THRESHOLD:
         return True
     # If only one test passed, check if it's a REPLAY_TEST
     return bool(pass_count == 1 and report[TestType.REPLAY_TEST]["passed"] == 1)

tests/test_critic.py

@@ -195,8 +195,7 @@ def test_generated_test_critic() -> None:
         timed_out=False,
         loop_index=1,
     )
-
-    test_results = [test_1, test_2, test_3, test_7]
+    test_results = [test_1, test_2, test_3, test_4, test_5, test_6, test_7, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -209,7 +208,7 @@ def test_generated_test_critic() -> None:
 
     assert quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_6, test_7]
+    test_results = [test_1, test_2, test_3, test_6, test_7, test_1, test_4, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -222,7 +221,7 @@ def test_generated_test_critic() -> None:
 
     assert quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_3, test_4, test_2, test_7]
+    test_results = [test_1, test_3, test_4, test_2, test_7, test_1, test_6, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -248,7 +247,7 @@ def test_generated_test_critic() -> None:
 
     assert not quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_4, test_5]
+    test_results = [test_1, test_2, test_3, test_4, test_5, test_1, test_1, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -287,7 +286,7 @@ def test_generated_test_critic() -> None:
 
     assert quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_4, test_5]
+    test_results = [test_1, test_2, test_3, test_4, test_5, test_1, test_1, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -328,7 +327,7 @@ def test_generated_test_critic() -> None:
 
     assert not quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_5]
+    test_results = [test_1, test_2, test_3, test_5, test_1, test_1, test_1, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,

tests/test_function_discovery.py

@@ -35,6 +35,17 @@ def test_function_eligible_for_optimization() -> None:
     assert len(functions_found[Path(f.name)]) == 0
 
 
+    # we want to trigger an error in the function discovery
+    function = """def test_invalid_code():"""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(function)
+        f.flush()
+        functions_found = find_all_functions_in_file(Path(f.name))
+    assert functions_found == {}
+
+
+
+
 def test_find_top_level_function_or_method():
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
         f.write(
@@ -82,6 +93,15 @@ def non_classmethod_function(cls, name):
         ).is_top_level
         # needed because this will be traced with a class_name being passed
 
+    # we want to write invalid code to ensure that the function discovery does not crash
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+            """def functionA():
+"""
+        )
+        f.flush()
+        path_obj_name = Path(f.name)
+        assert not inspect_top_level_functions_or_methods(path_obj_name, "functionA")
 
 def test_class_method_discovery():
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
@@ -152,6 +172,133 @@ def functionA():
             assert functions[file][0].function_name == "functionA"
 
 
+def test_nested_function():
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+"""
+import copy 
+
+def propagate_attributes(
+    nodes: dict[str, dict], edges: list[dict], source_node_id: str, attribute: str
+) -> dict[str, dict]:
+    modified_nodes = copy.deepcopy(nodes)
+
+    # Build an adjacency list for faster traversal
+    adjacency = {}
+    for edge in edges:
+        src = edge["source"]
+        tgt = edge["target"]
+        if src not in adjacency:
+            adjacency[src] = []
+        adjacency[src].append(tgt)
+
+    # Track visited nodes to avoid cycles
+    visited = set()
+
+    def traverse(node_id):
+        if node_id in visited:
+            return
+        visited.add(node_id)
+
+        # Propagate attribute from source node
+        if (
+            node_id != source_node_id
+            and source_node_id in modified_nodes
+            and attribute in modified_nodes[source_node_id]
+        ):
+            if node_id in modified_nodes:
+                modified_nodes[node_id][attribute] = modified_nodes[source_node_id][
+                    attribute
+                ]
+
+        # Continue propagation to neighbors
+        for neighbor in adjacency.get(node_id, []):
+            traverse(neighbor)
+
+    traverse(source_node_id)
+    return modified_nodes
+"""
+        )
+        f.flush()
+        test_config = TestConfig(
+            tests_root="tests", project_root_path=".", test_framework="pytest", tests_project_rootdir=Path()
+        )
+        path_obj_name = Path(f.name)
+        functions, functions_count = get_functions_to_optimize(
+            optimize_all=None,
+            replay_test=None,
+            file=path_obj_name,
+            test_cfg=test_config,
+            only_get_this_function=None,
+            ignore_paths=[Path("/bruh/")],
+            project_root=path_obj_name.parent,
+            module_root=path_obj_name.parent,
+        )
+
+        assert len(functions) == 1
+        assert functions_count == 1
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+"""
+def outer_function():
+    def inner_function():
+        pass
+
+    return inner_function
+"""
+        )
+        f.flush()
+        test_config = TestConfig(
+            tests_root="tests", project_root_path=".", test_framework="pytest", tests_project_rootdir=Path()
+        )
+        path_obj_name = Path(f.name)
+        functions, functions_count = get_functions_to_optimize(
+            optimize_all=None,
+            replay_test=None,
+            file=path_obj_name,
+            test_cfg=test_config,
+            only_get_this_function=None,
+            ignore_paths=[Path("/bruh/")],
+            project_root=path_obj_name.parent,
+            module_root=path_obj_name.parent,
+        )
+
+        assert len(functions) == 1
+        assert functions_count == 1
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+"""
+def outer_function():
+    def inner_function():
+        pass
+    
+    def another_inner_function():
+        pass
+    return inner_function, another_inner_function
+"""
+        )
+        f.flush()
+        test_config = TestConfig(
+            tests_root="tests", project_root_path=".", test_framework="pytest", tests_project_rootdir=Path()
+        )
+        path_obj_name = Path(f.name)
+        functions, functions_count = get_functions_to_optimize(
+            optimize_all=None,
+            replay_test=None,
+            file=path_obj_name,
+            test_cfg=test_config,
+            only_get_this_function=None,
+            ignore_paths=[Path("/bruh/")],
+            project_root=path_obj_name.parent,
+            module_root=path_obj_name.parent,
+        )
+
+        assert len(functions) == 1
+        assert functions_count == 1
+
+
 def test_filter_files_optimized():
     tests_root = Path("tests").resolve()
     module_root = Path().resolve()