Merge pull request #1457 from codeflash-ai/fix-coverage-qualified-name

Fix coverage function identification using qualified names

Author: KRRT7

codeflash/code_utils/concolic_utils.py

@@ -105,12 +105,12 @@ def clean_concolic_tests(test_suite_code: str) -> str:
         can_parse = False
         tree = None
 
-    if not can_parse:
+    if not can_parse or tree is None:
         return AssertCleanup().transform_asserts(test_suite_code)
 
     for node in ast.walk(tree):
         if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
-            new_body = []
+            new_body: list[ast.stmt] = []
             for stmt in node.body:
                 if isinstance(stmt, ast.Assert):
                     if isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Call):

codeflash/code_utils/coverage_utils.py

@@ -13,14 +13,28 @@
 def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]:
     """Extract the single dependent function from the code context excluding the main function."""
     dependent_functions = set()
+
+    # Compare using bare name since AST extracts bare function names
+    bare_main = main_function.rsplit(".", 1)[-1] if "." in main_function else main_function
+
     for code_string in code_context.testgen_context.code_strings:
-        ast_tree = ast.parse(code_string.code)
-        dependent_functions.update(
-            {node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))}
-        )
+        # Quick heuristic: skip parsing entirely if there is no 'def' token,
+        # since no function definitions can be present without it.
+        if "def" not in code_string.code:
+            continue
 
-    if main_function in dependent_functions:
-        dependent_functions.discard(main_function)
+        ast_tree = ast.parse(code_string.code)
+        # Add function names directly, skipping the bare main name.
+        for node in ast_tree.body:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                name = node.name
+                if name == bare_main:
+                    continue
+                dependent_functions.add(name)
+                # If more than one dependent function (other than the main) is found,
+                # we can return False early since the final result cannot be a single name.
+                if len(dependent_functions) > 1:
+                    return False
 
     if not dependent_functions:
         return False
@@ -32,6 +46,9 @@ def extract_dependent_function(main_function: str, code_context: CodeOptimizatio
 
 
 def build_fully_qualified_name(function_name: str, code_context: CodeOptimizationContext) -> str:
+    # If the name is already qualified (contains a dot), return as-is
+    if "." in function_name:
+        return function_name
     full_name = function_name
     for obj_name, parents in code_context.preexisting_objects:
         if obj_name == function_name:

codeflash/github/PrComment.py

@@ -26,10 +26,10 @@ class PrComment:
 
     def to_json(self) -> dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]]:
         report_table: dict[str, dict[str, int]] = {}
-        for test_type, result in self.winning_behavior_test_results.get_test_pass_fail_report_by_type().items():
+        for test_type, counts in self.winning_behavior_test_results.get_test_pass_fail_report_by_type().items():
             name = test_type.to_name()
             if name:
-                report_table[name] = result
+                report_table[name] = counts
 
         result: dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]] = {
             "optimization_explanation": self.optimization_explanation,

codeflash/optimization/function_optimizer.py

@@ -2788,7 +2788,7 @@ def run_and_parse_tests(
                 test_config=self.test_cfg,
                 optimization_iteration=optimization_iteration,
                 run_result=run_result,
-                function_name=self.function_to_optimize.function_name,
+                function_name=self.function_to_optimize.qualified_name,
                 source_file=self.function_to_optimize.file_path,
                 code_context=code_context,
                 coverage_database_file=coverage_database_file,

codeflash/tracing/pytest_parallelization.py

@@ -33,7 +33,7 @@ def pytest_split(
 
     except ImportError:
         return None, None
-    test_files = set()
+    test_files_set: set[str] = set()
 
     # Find all test_*.py files recursively in the directory
     for test_path in test_paths:
@@ -42,20 +42,20 @@ def pytest_split(
             return None, None
         if _test_path.is_dir():
             # Find all test files matching the pattern test_*.py
-            test_files.update(map(str, _test_path.rglob("test_*.py")))
-            test_files.update(map(str, _test_path.rglob("*_test.py")))
+            test_files_set.update(map(str, _test_path.rglob("test_*.py")))
+            test_files_set.update(map(str, _test_path.rglob("*_test.py")))
         elif _test_path.is_file():
-            test_files.add(str(_test_path))
+            test_files_set.add(str(_test_path))
 
-    if not test_files:
+    if not test_files_set:
         return [[]], None
 
     # Determine number of splits
     if num_splits is None:
         num_splits = os.cpu_count() or 4
 
     # randomize to increase chances of all splits being balanced
-    test_files = list(test_files)
+    test_files = list(test_files_set)
     shuffle(test_files)
 
     # Apply limit if specified
@@ -75,7 +75,7 @@ def pytest_split(
     chunk_size = ceil(total_files / num_splits)
 
     # Initialize result groups
-    result_groups = [[] for _ in range(num_splits)]
+    result_groups: list[list[str]] = [[] for _ in range(num_splits)]
 
     # Distribute files across groups
     for i, test_file in enumerate(test_files):

codeflash/verification/coverage_utils.py

@@ -351,7 +351,9 @@ def grab_dependent_function_from_coverage_data(
                 for file in files:
                     functions = files[file]["functions"]
                     for function in functions:
-                        if dependent_function_name in function:
+                        if function == dependent_function_name or (
+                            "." in dependent_function_name and function.endswith(f".{dependent_function_name}")
+                        ):
                             return FunctionCoverage(
                                 name=dependent_function_name,
                                 coverage=functions[function]["summary"]["percent_covered"],

tests/code_utils/test_coverage_utils.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+from typing import Any
+
+from codeflash.code_utils.coverage_utils import build_fully_qualified_name, extract_dependent_function
+from codeflash.models.function_types import FunctionParent
+from codeflash.models.models import CodeOptimizationContext, CodeString, CodeStringsMarkdown
+from codeflash.verification.coverage_utils import CoverageUtils
+
+
+def _make_code_context(
+    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
+    testgen_code_strings: list[CodeString] | None = None,
+) -> CodeOptimizationContext:
+    """Helper to create a minimal CodeOptimizationContext for testing."""
+    return CodeOptimizationContext(
+        testgen_context=CodeStringsMarkdown(code_strings=testgen_code_strings or []),
+        read_writable_code=CodeStringsMarkdown(),
+        helper_functions=[],
+        preexisting_objects=preexisting_objects,
+    )
+
+
+class TestBuildFullyQualifiedName:
+    def test_bare_name_with_class_parent(self) -> None:
+        ctx = _make_code_context({("__init__", (FunctionParent(name="HttpInterface", type="ClassDef"),))})
+        assert build_fully_qualified_name("__init__", ctx) == "HttpInterface.__init__"
+
+    def test_bare_name_no_parent(self) -> None:
+        ctx = _make_code_context({("helper_func", ())})
+        assert build_fully_qualified_name("helper_func", ctx) == "helper_func"
+
+    def test_already_qualified_name_returned_as_is(self) -> None:
+        """If name already contains a dot, skip preexisting_objects lookup."""
+        ctx = _make_code_context({("__init__", (FunctionParent(name="WrongClass", type="ClassDef"),))})
+        result = build_fully_qualified_name("HttpInterface.__init__", ctx)
+        assert result == "HttpInterface.__init__"
+
+    def test_bare_name_picks_first_match_from_set(self) -> None:
+        """With multiple __init__ entries, bare name picks an arbitrary one."""
+        ctx = _make_code_context(
+            {
+                ("__init__", (FunctionParent(name="ClassA", type="ClassDef"),)),
+                ("__init__", (FunctionParent(name="ClassB", type="ClassDef"),)),
+            }
+        )
+        result = build_fully_qualified_name("__init__", ctx)
+        assert result in {"ClassA.__init__", "ClassB.__init__"}
+
+    def test_qualified_name_avoids_ambiguity(self) -> None:
+        """Qualified name bypasses preexisting_objects entirely, avoiding ambiguity."""
+        ctx = _make_code_context(
+            {
+                ("__init__", (FunctionParent(name="ClassA", type="ClassDef"),)),
+                ("__init__", (FunctionParent(name="ClassB", type="ClassDef"),)),
+            }
+        )
+        assert build_fully_qualified_name("ClassB.__init__", ctx) == "ClassB.__init__"
+
+    def test_bare_name_not_in_preexisting_objects(self) -> None:
+        ctx = _make_code_context(set())
+        assert build_fully_qualified_name("some_func", ctx) == "some_func"
+
+    def test_nested_class_parent(self) -> None:
+        """Bare name under nested class parents gets fully qualified."""
+        ctx = _make_code_context(
+            {("method", (FunctionParent(name="Outer", type="ClassDef"), FunctionParent(name="Inner", type="ClassDef")))}
+        )
+        assert build_fully_qualified_name("method", ctx) == "Inner.Outer.method"
+
+    def test_non_classdef_parent_ignored(self) -> None:
+        """Only ClassDef parents are prepended to the name."""
+        ctx = _make_code_context({("helper", (FunctionParent(name="wrapper", type="FunctionDef"),))})
+        assert build_fully_qualified_name("helper", ctx) == "helper"
+
+
+class TestExtractDependentFunction:
+    def test_single_dependent_function(self) -> None:
+        ctx = _make_code_context(
+            preexisting_objects={("helper", ())},
+            testgen_code_strings=[CodeString(code="def main_func(): pass\ndef helper(): pass")],
+        )
+        result = extract_dependent_function("main_func", ctx)
+        assert result == "helper"
+
+    def test_qualified_main_function_discards_bare_match(self) -> None:
+        """Qualified main_function should still discard the matching bare name."""
+        ctx = _make_code_context(
+            preexisting_objects={("helper", ())},
+            testgen_code_strings=[CodeString(code="def __init__(): pass\ndef helper(): pass")],
+        )
+        result = extract_dependent_function("HttpInterface.__init__", ctx)
+        assert result == "helper"
+
+    def test_bare_main_function_discards_match(self) -> None:
+        """Bare main_function should still work for discarding."""
+        ctx = _make_code_context(
+            preexisting_objects={("helper", ())},
+            testgen_code_strings=[CodeString(code="def main_func(): pass\ndef helper(): pass")],
+        )
+        result = extract_dependent_function("main_func", ctx)
+        assert result == "helper"
+
+    def test_no_dependent_functions(self) -> None:
+        ctx = _make_code_context(preexisting_objects=set(), testgen_code_strings=[CodeString(code="x = 1\n")])
+        result = extract_dependent_function("main_func", ctx)
+        assert result is False
+
+    def test_multiple_dependent_functions_returns_false(self) -> None:
+        ctx = _make_code_context(
+            preexisting_objects=set(),
+            testgen_code_strings=[CodeString(code="def helper_a(): pass\ndef helper_b(): pass")],
+        )
+        result = extract_dependent_function("main_func", ctx)
+        assert result is False
+
+    def test_dependent_function_gets_qualified(self) -> None:
+        """The dependent function returned should be qualified via build_fully_qualified_name."""
+        ctx = _make_code_context(
+            preexisting_objects={("helper", (FunctionParent(name="MyClass", type="ClassDef"),))},
+            testgen_code_strings=[CodeString(code="def main_func(): pass\ndef helper(): pass")],
+        )
+        result = extract_dependent_function("main_func", ctx)
+        assert result == "MyClass.helper"
+
+    def test_only_main_in_code_returns_false(self) -> None:
+        """When code only contains the main function, no dependent function exists."""
+        ctx = _make_code_context(
+            preexisting_objects=set(), testgen_code_strings=[CodeString(code="def __init__(): pass")]
+        )
+        result = extract_dependent_function("HttpInterface.__init__", ctx)
+        assert result is False
+
+    def test_async_functions_extracted(self) -> None:
+        """Async function definitions are also extracted as dependent functions."""
+        ctx = _make_code_context(
+            preexisting_objects={("async_helper", ())},
+            testgen_code_strings=[CodeString(code="def main(): pass\nasync def async_helper(): pass")],
+        )
+        result = extract_dependent_function("main", ctx)
+        assert result == "async_helper"
+
+
+class TestGrabDependentFunctionFromCoverageData:
+    def _make_func_data(self, coverage_pct: float = 80.0) -> dict[str, Any]:
+        return {
+            "summary": {"percent_covered": coverage_pct},
+            "executed_lines": [1, 2, 3],
+            "missing_lines": [4],
+            "executed_branches": [[1, 0]],
+            "missing_branches": [[2, 1]],
+        }
+
+    def test_exact_match_in_coverage_data(self) -> None:
+        coverage_data = {"HttpInterface.__init__": self._make_func_data(90.0)}
+        result = CoverageUtils.grab_dependent_function_from_coverage_data("HttpInterface.__init__", coverage_data, {})
+        assert result.name == "HttpInterface.__init__"
+        assert result.coverage == 90.0
+
+    def test_fallback_exact_match_in_original_data(self) -> None:
+        original_cov_data = {
+            "files": {"http_api.py": {"functions": {"HttpInterface.__init__": self._make_func_data(75.0)}}}
+        }
+        result = CoverageUtils.grab_dependent_function_from_coverage_data(
+            "HttpInterface.__init__", {}, original_cov_data
+        )
+        assert result.name == "HttpInterface.__init__"
+        assert result.coverage == 75.0
+
+    def test_fallback_suffix_match_in_original_data(self) -> None:
+        """Qualified dependent name matches via suffix in original coverage data."""
+        original_cov_data = {
+            "files": {"http_api.py": {"functions": {"module.HttpInterface.__init__": self._make_func_data(60.0)}}}
+        }
+        result = CoverageUtils.grab_dependent_function_from_coverage_data(
+            "HttpInterface.__init__", {}, original_cov_data
+        )
+        assert result.name == "HttpInterface.__init__"
+        assert result.coverage == 60.0
+
+    def test_no_false_substring_match_bare_init(self) -> None:
+        """Bare __init__ should NOT match PathAwareCORSMiddleware.__init__ via substring."""
+        original_cov_data = {
+            "files": {"cors.py": {"functions": {"PathAwareCORSMiddleware.__init__": self._make_func_data(50.0)}}}
+        }
+        result = CoverageUtils.grab_dependent_function_from_coverage_data("__init__", {}, original_cov_data)
+        assert result.coverage == 0
+
+    def test_no_false_substring_match_different_class(self) -> None:
+        """Qualified name for one class should not match another class's method."""
+        original_cov_data = {
+            "files": {
+                "api.py": {
+                    "functions": {
+                        "PathAwareCORSMiddleware.__init__": self._make_func_data(50.0),
+                        "HttpInterface.__init__": self._make_func_data(85.0),
+                    }
+                }
+            }
+        }
+        result = CoverageUtils.grab_dependent_function_from_coverage_data(
+            "HttpInterface.__init__", {}, original_cov_data
+        )
+        assert result.name == "HttpInterface.__init__"
+        assert result.coverage == 85.0
+
+    def test_no_match_returns_zero_coverage(self) -> None:
+        result = CoverageUtils.grab_dependent_function_from_coverage_data("nonexistent_func", {}, {"files": {}})
+        assert result.coverage == 0
+        assert result.executed_lines == []
+
+    def test_qualified_suffix_no_match_for_partial_name(self) -> None:
+        """Ensure suffix match requires a dot boundary, not just string suffix."""
+        original_cov_data = {
+            "files": {"api.py": {"functions": {"XHttpInterface.__init__": self._make_func_data(40.0)}}}
+        }
+        # "HttpInterface.__init__" should NOT match "XHttpInterface.__init__" via suffix
+        result = CoverageUtils.grab_dependent_function_from_coverage_data(
+            "HttpInterface.__init__", {}, original_cov_data
+        )
+        assert result.coverage == 0
+
+    def test_bare_name_exact_match_in_fallback(self) -> None:
+        """Bare function name should still work with exact match in fallback."""
+        original_cov_data = {"files": {"utils.py": {"functions": {"helper_func": self._make_func_data(95.0)}}}}
+        result = CoverageUtils.grab_dependent_function_from_coverage_data("helper_func", {}, original_cov_data)
+        assert result.name == "helper_func"
+        assert result.coverage == 95.0