first pass

Author: KRRT7

codeflash/code_utils/code_replacer.py

@@ -336,3 +336,27 @@ def function_to_optimize_original_worktree_fqn(
         + "."
         + function_to_optimize.qualified_name
     )
+
+
+def clean_concolic_tests(test_suite_code: str) -> str:
+    try:
+        tree = ast.parse(test_suite_code)
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
+                new_body = []
+                for stmt in node.body:
+                    if isinstance(stmt, ast.Assert):
+                        if isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Call):
+                            new_body.append(ast.Expr(value=stmt.test.left))
+                        else:
+                            new_body.append(stmt)
+
+                    else:
+                        new_body.append(stmt)
+                node.body = new_body
+
+        return ast.unparse(tree).strip()
+    except SyntaxError:
+        logger.warning("Failed to parse and modify CrossHair generated tests.  Using original output.")
+        return test_suite_code

codeflash/verification/concolic_testing.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import ast
+import difflib
 import subprocess
 import tempfile
 from argparse import Namespace
 from pathlib import Path
 
 from codeflash.cli_cmds.console import console, logger
+from codeflash.code_utils.code_replacer import clean_concolic_tests
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
 from codeflash.code_utils.static_analysis import has_typed_parameters
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
@@ -21,7 +23,11 @@ def generate_concolic_tests(
 ) -> tuple[dict[str, list[FunctionCalledInTest]], str]:
     function_to_concolic_tests = {}
     concolic_test_suite_code = ""
-    if test_cfg.concolic_test_root_dir and has_typed_parameters(function_to_optimize_ast, function_to_optimize.parents):
+    if (
+        test_cfg.concolic_test_root_dir
+        and isinstance(function_to_optimize_ast, (ast.FunctionDef, ast.AsyncFunctionDef))
+        and has_typed_parameters(function_to_optimize_ast, function_to_optimize.parents)
+    ):
         logger.info("Generating concolic opcode coverage tests for the original code…")
         console.rule()
         try:
@@ -54,7 +60,8 @@ def generate_concolic_tests(
             return function_to_concolic_tests, concolic_test_suite_code
 
         if cover_result.returncode == 0:
-            concolic_test_suite_code: str = cover_result.stdout
+            original_code: str = cover_result.stdout
+            concolic_test_suite_code: str = clean_concolic_tests(original_code)
             concolic_test_suite_dir = Path(tempfile.mkdtemp(dir=test_cfg.concolic_test_root_dir))
             concolic_test_suite_path = concolic_test_suite_dir / "test_concolic_coverage.py"
             concolic_test_suite_path.write_text(concolic_test_suite_code, encoding="utf8")

tests/test_code_utils.py

@@ -6,6 +6,7 @@
 
 import pytest
 
+from codeflash.code_utils.code_replacer import clean_concolic_tests
 from codeflash.code_utils.code_utils import (
     cleanup_paths,
     file_name_from_test_module_name,
@@ -378,3 +379,51 @@ def test_prepare_coverage_files(mock_get_run_tmp_file: MagicMock) -> None:
     assert coverage_database_file == mock_coverage_file
     assert coveragercfile == mock_coveragerc_file
     mock_coveragerc_file.write_text.assert_called_once_with(f"[run]\n branch = True\ndata_file={mock_coverage_file}\n")
+
+
+def test_clean_concolic_tests() -> None:
+    original_code = """
+def test_add_numbers(x: int, y: int) -> None:
+    assert add_numbers(1, 2) == 3
+
+
+def test_concatenate_strings(s1: str, s2: str) -> None:
+    assert concatenate_strings("hello", "world") == "helloworld"
+
+
+def test_append_to_list(my_list: list[int], element: int) -> None:
+    assert append_to_list([1, 2, 3], 4) == [1, 2, 3, 4]
+
+
+def test_get_dict_value(my_dict: dict[str, int], key: str) -> None:
+    assert get_dict_value({"a": 1, "b": 2}, "a") == 1
+
+
+def test_union_sets(set1: set[int], set2: set[int]) -> None:
+    assert union_sets({1, 2, 3}, {3, 4, 5}) == {1, 2, 3, 4, 5}
+
+def test_calculate_tuple_sum(my_tuple: tuple[int, int, int]) -> None:
+    assert calculate_tuple_sum((1, 2, 3)) == 6
+"""
+
+    cleaned_code = clean_concolic_tests(original_code)
+    expected_cleaned_code = """
+def test_add_numbers(x: int, y: int) -> None:
+    add_numbers(1, 2)
+
+def test_concatenate_strings(s1: str, s2: str) -> None:
+    concatenate_strings('hello', 'world')
+
+def test_append_to_list(my_list: list[int], element: int) -> None:
+    append_to_list([1, 2, 3], 4)
+
+def test_get_dict_value(my_dict: dict[str, int], key: str) -> None:
+    get_dict_value({'a': 1, 'b': 2}, 'a')
+
+def test_union_sets(set1: set[int], set2: set[int]) -> None:
+    union_sets({1, 2, 3}, {3, 4, 5})
+
+def test_calculate_tuple_sum(my_tuple: tuple[int, int, int]) -> None:
+    calculate_tuple_sum((1, 2, 3))
+"""
+    assert cleaned_code == expected_cleaned_code.strip()