Merge pull request #1663 from codeflash-ai/fix/java-maven-test-execution-bugs

fix: resolve Maven test execution blockers for open-source Java repos

Author: mohamedashrraf222

codeflash/languages/java/instrumentation.py

@@ -1325,32 +1325,3 @@ def instrument_generated_java_test(
 
     logger.debug("Instrumented generated Java test for %s (mode=%s)", function_name, mode)
     return modified_code
-
-
-def _add_import(source: str, import_statement: str) -> str:
-    """Add an import statement to the source.
-
-    Args:
-        source: The source code.
-        import_statement: The import to add.
-
-    Returns:
-        Source with import added.
-
-    """
-    lines = source.splitlines(keepends=True)
-    insert_idx = 0
-
-    # Find the last import or package statement
-    for i, line in enumerate(lines):
-        stripped = line.strip()
-        if stripped.startswith(("import ", "package ")):
-            insert_idx = i + 1
-        elif stripped and not stripped.startswith("//") and not stripped.startswith("/*"):
-            # First non-import, non-comment line
-            if insert_idx == 0:
-                insert_idx = i
-            break
-
-    lines.insert(insert_idx, import_statement + "\n")
-    return "".join(lines)

codeflash/languages/java/remove_asserts.py

@@ -198,6 +198,14 @@ def __init__(
         # Precompile regex to find next special character (quotes, parens, braces).
         self._special_re = re.compile(r"[\"'{}()]")
 
+        # Precompile literal/cast regexes to avoid recompilation on each literal check.
+        self._LONG_LITERAL_RE = re.compile(r"^-?\d+[lL]$")
+        self._INT_LITERAL_RE = re.compile(r"^-?\d+$")
+        self._DOUBLE_LITERAL_RE = re.compile(r"^-?\d+\.\d*[dD]?$|^-?\d+[dD]$")
+        self._FLOAT_LITERAL_RE = re.compile(r"^-?\d+\.?\d*[fF]$")
+        self._CHAR_LITERAL_RE = re.compile(r"^'.'$|^'\\.'$")
+        self._cast_re = re.compile(r"^\((\w+)\)")
+
     def transform(self, source: str) -> str:
         """Remove assertions from source code, preserving target function calls.
 
@@ -894,6 +902,143 @@ def _find_balanced_braces(self, code: str, open_brace_pos: int) -> tuple[str | N
 
         return code[open_brace_pos + 1 : pos - 1], pos
 
+    def _infer_return_type(self, assertion: AssertionMatch) -> str:
+        """Infer the Java return type from the assertion context.
+
+        For assertEquals(expected, actual) patterns, the expected literal determines the type.
+        For assertTrue/assertFalse, the result is boolean.
+        Falls back to Object when the type cannot be determined.
+        """
+        method = assertion.assertion_method
+
+        # assertTrue/assertFalse always deal with boolean values
+        if method in {"assertTrue", "assertFalse"}:
+            return "boolean"
+
+        # assertNull/assertNotNull — keep Object (reference type)
+        if method in {"assertNull", "assertNotNull"}:
+            return "Object"
+
+        # For assertEquals/assertNotEquals/assertSame, try to infer from the expected literal
+        if method in JUNIT5_VALUE_ASSERTIONS:
+            return self._infer_type_from_assertion_args(assertion.original_text, method)
+
+        # For fluent assertions (assertThat), type inference is harder — keep Object
+        return "Object"
+
+    # Regex patterns for Java literal type inference
+    _LONG_LITERAL_RE = re.compile(r"^-?\d+[lL]$")
+    _INT_LITERAL_RE = re.compile(r"^-?\d+$")
+    _DOUBLE_LITERAL_RE = re.compile(r"^-?\d+\.\d*[dD]?$|^-?\d+[dD]$")
+    _FLOAT_LITERAL_RE = re.compile(r"^-?\d+\.?\d*[fF]$")
+    _CHAR_LITERAL_RE = re.compile(r"^'.'$|^'\\.'$")
+
+    def _infer_type_from_assertion_args(self, original_text: str, method: str) -> str:
+        """Infer the return type from assertEquals/assertNotEquals expected value."""
+        # Extract the args portion from the assertion text
+        # Pattern: assertXxx( args... )
+        paren_idx = original_text.find("(")
+        if paren_idx < 0:
+            return "Object"
+
+        args_str = original_text[paren_idx + 1 :]
+        # Remove trailing ");", whitespace
+        args_str = args_str.rstrip()
+        if args_str.endswith(");"):
+            args_str = args_str[:-2]
+        elif args_str.endswith(")"):
+            args_str = args_str[:-1]
+
+        # Fast-path: only extract the first top-level argument instead of splitting all arguments.
+        first_arg = self._extract_first_arg(args_str)
+        if not first_arg:
+            return "Object"
+
+        expected = first_arg.strip()
+
+        # JUnit 4 has assertEquals(String message, expected, actual) where the first arg is a message.
+        # If the first arg is a string literal, check if there are 3+ args — if so, the real expected
+        # value is the second argument, not the message string.
+        if expected.startswith('"') and method in ("assertEquals", "assertNotEquals"):
+            all_args = self._split_top_level_args(args_str)
+            if len(all_args) >= 3:
+                expected = all_args[1].strip()
+
+        return self._type_from_literal(expected)
+
+    def _type_from_literal(self, value: str) -> str:
+        """Determine the Java type of a literal value."""
+        if value in ("true", "false"):
+            return "boolean"
+        if value == "null":
+            return "Object"
+        if self._FLOAT_LITERAL_RE.match(value):
+            return "float"
+        if self._DOUBLE_LITERAL_RE.match(value):
+            return "double"
+        if self._LONG_LITERAL_RE.match(value):
+            return "long"
+        if self._INT_LITERAL_RE.match(value):
+            return "int"
+        if self._CHAR_LITERAL_RE.match(value):
+            return "char"
+        if value.startswith('"'):
+            return "String"
+        # Cast expression like (byte)0, (short)1
+        cast_match = self._cast_re.match(value)
+        if cast_match:
+            return cast_match.group(1)
+        return "Object"
+
+    def _split_top_level_args(self, args_str: str) -> list[str]:
+        """Split assertion arguments at top-level commas, respecting parens/strings/generics."""
+        # Fast-path: if there are no special delimiters that require parsing,
+        # we can use a simple split which is much faster for common simple cases.
+        if not self._special_re.search(args_str):
+            # Preserve original behavior of returning a list with the single unstripped string
+            # when there are no commas, otherwise split on commas.
+            if "," in args_str:
+                return args_str.split(",")
+            return [args_str]
+
+        args: list[str] = []
+        depth = 0
+        current: list[str] = []
+        i = 0
+        in_string = False
+        string_char = ""
+
+        while i < len(args_str):
+            ch = args_str[i]
+
+            if in_string:
+                current.append(ch)
+                if ch == "\\" and i + 1 < len(args_str):
+                    i += 1
+                    current.append(args_str[i])
+                elif ch == string_char:
+                    in_string = False
+            elif ch in ('"', "'"):
+                in_string = True
+                string_char = ch
+                current.append(ch)
+            elif ch in ("(", "<", "[", "{"):
+                depth += 1
+                current.append(ch)
+            elif ch in (")", ">", "]", "}"):
+                depth -= 1
+                current.append(ch)
+            elif ch == "," and depth == 0:
+                args.append("".join(current))
+                current = []
+            else:
+                current.append(ch)
+            i += 1
+
+        if current:
+            args.append("".join(current))
+        return args
+
     def _generate_replacement(self, assertion: AssertionMatch) -> str:
         """Generate replacement code for an assertion.
 
@@ -912,18 +1057,34 @@ def _generate_replacement(self, assertion: AssertionMatch) -> str:
         if not assertion.target_calls:
             return ""
 
+        # Infer the return type from assertion context to avoid Object→primitive cast errors
+        return_type = self._infer_return_type(assertion)
+
         # Generate capture statements for each target call
-        replacements = []
+        replacements: list[str] = []
         # For the first replacement, use the full leading whitespace
         # For subsequent ones, strip leading newlines to avoid extra blank lines
-        base_indent = assertion.leading_whitespace.lstrip("\n\r")
-        for i, call in enumerate(assertion.target_calls):
-            self.invocation_counter += 1
-            var_name = f"_cf_result{self.invocation_counter}"
-            if i == 0:
-                replacements.append(f"{assertion.leading_whitespace}Object {var_name} = {call.full_call};")
-            else:
-                replacements.append(f"{base_indent}Object {var_name} = {call.full_call};")
+        leading_ws = assertion.leading_whitespace
+        base_indent = leading_ws.lstrip("\n\r")
+
+        # Use a local counter to minimize attribute write overhead in the loop.
+        inv = self.invocation_counter
+
+        calls = assertion.target_calls
+        # Handle first call explicitly to avoid a per-iteration branch
+        if calls:
+            inv += 1
+            var_name = "_cf_result" + str(inv)
+            replacements.append(f"{leading_ws}{return_type} {var_name} = {calls[0].full_call};")
+
+            # Handle remaining calls
+            for call in calls[1:]:
+                inv += 1
+                var_name = "_cf_result" + str(inv)
+                replacements.append(f"{base_indent}{return_type} {var_name} = {call.full_call};")
+
+        # Write back the counter
+        self.invocation_counter = inv
 
         return "\n".join(replacements)
 
@@ -942,8 +1103,10 @@ def _generate_exception_replacement(self, assertion: AssertionMatch) -> str:
             try { code(); } catch (IllegalArgumentException _cf_caught1) { ex = _cf_caught1; } catch (Exception _cf_ignored1) {}
 
         """
-        self.invocation_counter += 1
-        counter = self.invocation_counter
+        # Increment invocation counter once for this exception handling
+        inv = self.invocation_counter + 1
+        self.invocation_counter = inv
+        counter = inv
         ws = assertion.leading_whitespace
         base_indent = ws.lstrip("\n\r")
 
@@ -982,6 +1145,58 @@ def _generate_exception_replacement(self, assertion: AssertionMatch) -> str:
         # Fallback: comment out the assertion
         return f"{ws}// Removed assertThrows: could not extract callable"
 
+    def _extract_first_arg(self, args_str: str) -> str | None:
+        """Extract the first top-level argument from args_str.
+
+        This is a lightweight alternative to splitting all top-level arguments;
+        it stops at the first top-level comma, respects nested delimiters and strings,
+        and avoids constructing the full argument list for better performance.
+        """
+        n = len(args_str)
+        i = 0
+
+        # skip leading whitespace
+        while i < n and args_str[i].isspace():
+            i += 1
+        if i >= n:
+            return None
+
+        depth = 0
+        in_string = False
+        string_char = ""
+        cur: list[str] = []
+
+        while i < n:
+            ch = args_str[i]
+
+            if in_string:
+                cur.append(ch)
+                if ch == "\\" and i + 1 < n:
+                    i += 1
+                    cur.append(args_str[i])
+                elif ch == string_char:
+                    in_string = False
+            elif ch in ('"', "'"):
+                in_string = True
+                string_char = ch
+                cur.append(ch)
+            elif ch in ("(", "<", "[", "{"):
+                depth += 1
+                cur.append(ch)
+            elif ch in (")", ">", "]", "}"):
+                depth -= 1
+                cur.append(ch)
+            elif ch == "," and depth == 0:
+                break
+            else:
+                cur.append(ch)
+            i += 1
+
+        # Trim trailing whitespace from the extracted argument
+        if not cur:
+            return None
+        return "".join(cur).rstrip()
+
 
 def transform_java_assertions(source: str, function_name: str, qualified_name: str | None = None) -> str:
     """Transform Java test code by removing assertions and capturing function calls.

codeflash/languages/java/test_runner.py

@@ -47,6 +47,17 @@
 # Allows: letters, digits, underscores, dots, and dollar signs (inner classes)
 _VALID_JAVA_CLASS_NAME = re.compile(r"^[a-zA-Z_$][a-zA-Z0-9_$.]*$")
 
+# Skip validation/analysis plugins that reject generated instrumented files
+# (e.g. Apache Rat rejects missing license headers, Checkstyle rejects naming, etc.)
+_MAVEN_VALIDATION_SKIP_FLAGS = [
+    "-Drat.skip=true",
+    "-Dcheckstyle.skip=true",
+    "-Dspotbugs.skip=true",
+    "-Dpmd.skip=true",
+    "-Denforcer.skip=true",
+    "-Djapicmp.skip=true",
+]
+
 
 def _run_cmd_kill_pg_on_timeout(
     cmd: list[str],
@@ -85,9 +96,7 @@ def _run_cmd_kill_pg_on_timeout(
         # Windows does not have POSIX process groups / killpg.  Fall back to
         # the standard subprocess.run() behaviour (kills parent only).
         try:
-            return subprocess.run(
-                cmd, cwd=cwd, env=env, capture_output=True, text=text, timeout=timeout, check=False
-            )
+            return subprocess.run(cmd, cwd=cwd, env=env, capture_output=True, text=text, timeout=timeout, check=False)
         except subprocess.TimeoutExpired:
             return subprocess.CompletedProcess(
                 args=cmd, returncode=-2, stdout="", stderr=f"Process timed out after {timeout}s"
@@ -509,8 +518,9 @@ def run_behavioral_tests(
                     add_jacoco_plugin_to_pom(pom_path)
                 coverage_xml_path = get_jacoco_xml_path(project_root)
 
-    # Use a minimum timeout of 60s for Java builds (120s when coverage is enabled due to verify phase)
-    min_timeout = 120 if enable_coverage else 60
+    # Use a minimum timeout of 60s for Java builds (300s when coverage is enabled due to verify phase
+    # which runs full compilation + instrumentation + test execution in multi-module projects)
+    min_timeout = 300 if enable_coverage else 60
     effective_timeout = max(timeout or 300, min_timeout)
 
     if enable_coverage:
@@ -591,6 +601,7 @@ def _compile_tests(
         return subprocess.CompletedProcess(args=["mvn"], returncode=-1, stdout="", stderr="Maven not found")
 
     cmd = [mvn, "test-compile", "-e", "-B"]  # Show errors but not verbose output; -B for batch mode (no ANSI colors)
+    cmd.extend(_MAVEN_VALIDATION_SKIP_FLAGS)
 
     if test_module:
         cmd.extend(["-pl", test_module, "-am"])
@@ -1526,6 +1537,7 @@ def _run_maven_tests(
     # JaCoCo's report goal is bound to the verify phase to get post-test execution data
     maven_goal = "verify" if enable_coverage else "test"
     cmd = [mvn, maven_goal, "-fae", "-B"]  # Fail at end to run all tests; -B for batch mode (no ANSI colors)
+    cmd.extend(_MAVEN_VALIDATION_SKIP_FLAGS)
 
     # Add --add-opens flags for Java 16+ module system compatibility.
     # The codeflash-runtime Serializer uses Kryo which needs reflective access to
@@ -1562,7 +1574,16 @@ def _run_maven_tests(
         # -am = also make dependencies
         # -DfailIfNoTests=false allows dependency modules without tests to pass
         # -DskipTests=false overrides any skipTests=true in pom.xml
-        cmd.extend(["-pl", test_module, "-am", "-DfailIfNoTests=false", "-DskipTests=false"])
+        cmd.extend(
+            [
+                "-pl",
+                test_module,
+                "-am",
+                "-DfailIfNoTests=false",
+                "-Dsurefire.failIfNoSpecifiedTests=false",
+                "-DskipTests=false",
+            ]
+        )
 
     if test_filter:
         # Validate test filter to prevent command injection

codeflash/verification/parse_test_output.py

@@ -786,7 +786,13 @@ def parse_test_xml(
             if class_name is not None and class_name.startswith(test_module_path):
                 test_class = class_name[len(test_module_path) + 1 :]  # +1 for the dot, gets Unittest class name
 
-            loop_index = int(testcase.name.split("[ ")[-1][:-2]) if testcase.name and "[" in testcase.name else 1
+            loop_index = 1
+            if testcase.name and "[" in testcase.name:
+                bracket_content = testcase.name.rsplit("[", 1)[-1].rstrip("]").strip()
+                try:
+                    loop_index = int(bracket_content)
+                except ValueError:
+                    loop_index = 1
 
             timed_out = False
             if len(testcase.result) > 1: