From a65f1d1928b36409eb269c94e7f47976f4bdf159 Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Wed, 25 Feb 2026 15:44:42 +0000
Subject: [PATCH 01/11] fix: skip Maven validation plugins that reject
 generated test files

Maven plugins like Apache Rat, Checkstyle, SpotBugs, PMD, Enforcer, and
japicmp reject generated instrumented Java files (e.g. missing license
headers). Skip these validation plugins during test compilation and
execution since they are irrelevant for generated test code.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 codeflash/languages/java/test_runner.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/codeflash/languages/java/test_runner.py b/codeflash/languages/java/test_runner.py
index 45a23b9e7..146453a07 100644
--- a/codeflash/languages/java/test_runner.py
+++ b/codeflash/languages/java/test_runner.py
@@ -45,6 +45,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 _validate_java_class_name(class_name: str) -> bool:
     """Validate that a string is a valid Java class name.
@@ -493,6 +504,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"])
@@ -1447,6 +1459,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

From 29dd7a384f92401489bae38fac272f9ffef79136 Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Wed, 25 Feb 2026 15:44:55 +0000
Subject: [PATCH 02/11] fix: add surefire.failIfNoSpecifiedTests for
 multi-module Maven builds

In multi-module projects like Guava, -Dtest=X filter matches zero tests
in dependency modules built with -am, causing "No tests matching pattern"
failures. Adding -Dsurefire.failIfNoSpecifiedTests=false allows modules
with no matching tests to pass while still running the correct tests in
the target module.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 codeflash/languages/java/test_runner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/codeflash/languages/java/test_runner.py b/codeflash/languages/java/test_runner.py
index 146453a07..8dcaf401a 100644
--- a/codeflash/languages/java/test_runner.py
+++ b/codeflash/languages/java/test_runner.py
@@ -1496,7 +1496,7 @@ 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

From 15ade968dd811089579f9ab9acfd5ace2db2a330 Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Wed, 25 Feb 2026 15:45:16 +0000
Subject: [PATCH 03/11] fix: auto-add missing standard library imports in
 AI-generated Java tests

AI-generated test code sometimes uses standard library classes like
Arrays, List, HashMap etc. without the corresponding import statement,
causing compilation failures. Add ensure_common_java_imports() that
detects usage of common classes and adds missing imports automatically.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 codeflash/languages/java/instrumentation.py | 30 +++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/codeflash/languages/java/instrumentation.py b/codeflash/languages/java/instrumentation.py
index 4a2a0e0a9..81e125ae6 100644
--- a/codeflash/languages/java/instrumentation.py
+++ b/codeflash/languages/java/instrumentation.py
@@ -1260,6 +1260,35 @@ def remove_instrumentation(source: str) -> str:
     return source
 
 
+_COMMON_JAVA_IMPORTS = {
+    "Arrays": "import java.util.Arrays;",
+    "List": "import java.util.List;",
+    "ArrayList": "import java.util.ArrayList;",
+    "Map": "import java.util.Map;",
+    "HashMap": "import java.util.HashMap;",
+    "Set": "import java.util.Set;",
+    "HashSet": "import java.util.HashSet;",
+    "Collections": "import java.util.Collections;",
+    "Collectors": "import java.util.stream.Collectors;",
+    "Random": "import java.util.Random;",
+    "BigDecimal": "import java.math.BigDecimal;",
+    "BigInteger": "import java.math.BigInteger;",
+}
+
+
+def ensure_common_java_imports(test_code: str) -> str:
+    for class_name, import_stmt in _COMMON_JAVA_IMPORTS.items():
+        if not re.search(rf"\b{class_name}\b", test_code):
+            continue
+        if import_stmt in test_code:
+            continue
+        package = import_stmt.split()[1].rsplit(".", 1)[0]
+        if f"import {package}.*;" in test_code:
+            continue
+        test_code = _add_import(test_code, import_stmt)
+    return test_code
+
+
 def instrument_generated_java_test(
     test_code: str,
     function_name: str,
@@ -1290,6 +1319,7 @@ def instrument_generated_java_test(
     from codeflash.languages.java.remove_asserts import transform_java_assertions
 
     test_code = transform_java_assertions(test_code, function_name, qualified_name)
+    test_code = ensure_common_java_imports(test_code)
 
     # Extract class name from the test code
     # Use pattern that starts at beginning of line to avoid matching words in comments

From 4549220c9ce8f2063baa5bfba351165884f759dd Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Wed, 25 Feb 2026 20:19:13 +0000
Subject: [PATCH 04/11] fix: infer Java return types in assertion transformer
 instead of using Object

The assertion transformer always declared `Object _cf_resultN = call()` when
replacing assertions, losing the actual return type. This caused compilation
failures when the result was used in a context expecting a primitive type
(e.g., int, boolean).

Now infers the return type from assertion context:
- assertEquals(int_literal, call()) -> int
- assertTrue/assertFalse(call()) -> boolean
- assertEquals("string", call()) -> String
- Falls back to Object when type can't be determined

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 codeflash/languages/java/remove_asserts.py    | 130 +++++++++++++++++-
 .../test_java/test_remove_asserts.py          | 106 +++++++-------
 2 files changed, 181 insertions(+), 55 deletions(-)

diff --git a/codeflash/languages/java/remove_asserts.py b/codeflash/languages/java/remove_asserts.py
index 8544a771d..d89e76cce 100644
--- a/codeflash/languages/java/remove_asserts.py
+++ b/codeflash/languages/java/remove_asserts.py
@@ -894,6 +894,129 @@ 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]
+
+        # Split top-level args (respecting parens, strings, generics)
+        args = self._split_top_level_args(args_str)
+        if not args:
+            return "Object"
+
+        # assertEquals has (expected, actual) or (expected, actual, message/delta)
+        # Some overloads have (message, expected, actual) in JUnit 4 but JUnit 5 uses (expected, actual[, message])
+        # Try the first argument as the expected value
+        expected = args[0].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 = re.match(r"^\((\w+)\)", 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."""
+        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,6 +1035,9 @@ 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 = []
         # For the first replacement, use the full leading whitespace
@@ -921,9 +1047,9 @@ def _generate_replacement(self, assertion: AssertionMatch) -> str:
             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};")
+                replacements.append(f"{assertion.leading_whitespace}{return_type} {var_name} = {call.full_call};")
             else:
-                replacements.append(f"{base_indent}Object {var_name} = {call.full_call};")
+                replacements.append(f"{base_indent}{return_type} {var_name} = {call.full_call};")
 
         return "\n".join(replacements)
 
diff --git a/tests/test_languages/test_java/test_remove_asserts.py b/tests/test_languages/test_java/test_remove_asserts.py
index e0a252ad8..edc7138ce 100644
--- a/tests/test_languages/test_java/test_remove_asserts.py
+++ b/tests/test_languages/test_java/test_remove_asserts.py
@@ -41,7 +41,7 @@ def test_assertfalse_with_message(self):
 public class BitSetTest {
     @Test
     public void testGet_IndexZero_ReturnsFalse() {
-        Object _cf_result1 = instance.get(0);
+        boolean _cf_result1 = instance.get(0);
     }
 }
 """
@@ -67,7 +67,7 @@ def test_asserttrue_with_message(self):
 public class BitSetTest {
     @Test
     public void testGet_SetBit_DetectedTrue() {
-        Object _cf_result1 = bs.get(67);
+        boolean _cf_result1 = bs.get(67);
     }
 }
 """
@@ -93,7 +93,7 @@ def test_assertequals_with_static_call(self):
 public class FibonacciTest {
     @Test
     public void testFibonacci() {
-        Object _cf_result1 = Fibonacci.fibonacci(10);
+        int _cf_result1 = Fibonacci.fibonacci(10);
     }
 }
 """
@@ -121,7 +121,7 @@ def test_assertequals_with_instance_call(self):
     @Test
     public void testAdd() {
         Calculator calc = new Calculator();
-        Object _cf_result1 = calc.add(2, 2);
+        int _cf_result1 = calc.add(2, 2);
     }
 }
 """
@@ -199,7 +199,7 @@ def test_assertnotequals(self):
 public class CalculatorTest {
     @Test
     public void testSubtract() {
-        Object _cf_result1 = calc.subtract(5, 3);
+        int _cf_result1 = calc.subtract(5, 3);
     }
 }
 """
@@ -251,7 +251,7 @@ def test_qualified_assert_call(self):
 public class CalculatorTest {
     @Test
     public void testAdd() {
-        Object _cf_result1 = calc.add(2, 2);
+        int _cf_result1 = calc.add(2, 2);
     }
 }
 """
@@ -298,9 +298,9 @@ def test_assertequals_static_import(self):
 public class FibonacciTest {
     @Test
     void testFibonacci() {
-        Object _cf_result1 = Fibonacci.fibonacci(0);
-        Object _cf_result2 = Fibonacci.fibonacci(1);
-        Object _cf_result3 = Fibonacci.fibonacci(10);
+        int _cf_result1 = Fibonacci.fibonacci(0);
+        int _cf_result2 = Fibonacci.fibonacci(1);
+        int _cf_result3 = Fibonacci.fibonacci(10);
     }
 }
 """
@@ -326,7 +326,7 @@ def test_assertequals_qualified(self):
 public class FibonacciTest {
     @Test
     void testFibonacci() {
-        Object _cf_result1 = Fibonacci.fibonacci(10);
+        int _cf_result1 = Fibonacci.fibonacci(10);
     }
 }
 """
@@ -485,7 +485,7 @@ def test_asserttrue_boolean_call(self):
 public class FibonacciTest {
     @Test
     void testIsFibonacci() {
-        Object _cf_result1 = Fibonacci.isFibonacci(5);
+        boolean _cf_result1 = Fibonacci.isFibonacci(5);
     }
 }
 """
@@ -511,7 +511,7 @@ def test_assertfalse_boolean_call(self):
 public class FibonacciTest {
     @Test
     void testIsNotFibonacci() {
-        Object _cf_result1 = Fibonacci.isFibonacci(4);
+        boolean _cf_result1 = Fibonacci.isFibonacci(4);
     }
 }
 """
@@ -709,7 +709,7 @@ def test_multiple_calls_in_one_assertion(self):
 public class FibonacciTest {
     @Test
     void testConsecutive() {
-        Object _cf_result1 = Fibonacci.areConsecutiveFibonacci(Fibonacci.fibonacci(5), Fibonacci.fibonacci(6));
+        boolean _cf_result1 = Fibonacci.areConsecutiveFibonacci(Fibonacci.fibonacci(5), Fibonacci.fibonacci(6));
     }
 }
 """
@@ -739,11 +739,11 @@ def test_multiple_assertions_in_one_method(self):
 public class FibonacciTest {
     @Test
     void testMultiple() {
-        Object _cf_result1 = Fibonacci.fibonacci(0);
-        Object _cf_result2 = Fibonacci.fibonacci(1);
-        Object _cf_result3 = Fibonacci.fibonacci(2);
-        Object _cf_result4 = Fibonacci.fibonacci(3);
-        Object _cf_result5 = Fibonacci.fibonacci(5);
+        int _cf_result1 = Fibonacci.fibonacci(0);
+        int _cf_result2 = Fibonacci.fibonacci(1);
+        int _cf_result3 = Fibonacci.fibonacci(2);
+        int _cf_result4 = Fibonacci.fibonacci(3);
+        int _cf_result5 = Fibonacci.fibonacci(5);
     }
 }
 """
@@ -774,7 +774,7 @@ def test_assertion_without_target_removed(self):
 public class SetupTest {
     @Test
     void testSetup() {
-        Object _cf_result1 = Fibonacci.fibonacci(10);
+        int _cf_result1 = Fibonacci.fibonacci(10);
     }
 }
 """
@@ -829,7 +829,7 @@ def test_multiline_assertion(self):
 public class FibonacciTest {
     @Test
     void testFibonacci() {
-        Object _cf_result1 = Fibonacci.fibonacci(10);
+        int _cf_result1 = Fibonacci.fibonacci(10);
     }
 }
 """
@@ -855,7 +855,7 @@ def test_assertion_with_string_containing_parens(self):
 public class ParserTest {
     @Test
     void testParse() {
-        Object _cf_result1 = parser.parse("input(1)");
+        String _cf_result1 = parser.parse("input(1)");
     }
 }
 """
@@ -911,7 +911,7 @@ def test_nested_method_calls(self):
 public class FibonacciTest {
     @Test
     void testIndex() {
-        Object _cf_result1 = Fibonacci.fibonacciIndex(Fibonacci.fibonacci(10));
+        int _cf_result1 = Fibonacci.fibonacciIndex(Fibonacci.fibonacci(10));
     }
 }
 """
@@ -937,7 +937,7 @@ def test_chained_method_on_result(self):
 public class FibonacciTest {
     @Test
     void testUpTo() {
-        Object _cf_result1 = Fibonacci.fibonacciUpTo(20);
+        int _cf_result1 = Fibonacci.fibonacciUpTo(20);
     }
 }
 """
@@ -1053,24 +1053,24 @@ class TestBitSetLikeQuestDB:
 
     @Test
     public void testGet_IndexZero_ReturnsFalse() {
-        Object _cf_result1 = instance.get(0);
+        boolean _cf_result1 = instance.get(0);
     }
 
     @Test
     public void testGet_SpecificIndexWithinRange_ReturnsFalse() {
-        Object _cf_result2 = instance.get(100);
+        boolean _cf_result2 = instance.get(100);
     }
 
     @Test
     public void testGet_LastIndexOfInitialRange_ReturnsFalse() {
         int lastIndex = 16 * BitSet.BITS_PER_WORD - 1;
-        Object _cf_result3 = instance.get(lastIndex);
+        boolean _cf_result3 = instance.get(lastIndex);
     }
 
     @Test
     public void testGet_IndexBeyondAllocated_ReturnsFalse() {
         int beyond = 16 * BitSet.BITS_PER_WORD;
-        Object _cf_result4 = instance.get(beyond);
+        boolean _cf_result4 = instance.get(beyond);
     }
 
     @Test(expected = ArrayIndexOutOfBoundsException.class)
@@ -1086,22 +1086,22 @@ class TestBitSetLikeQuestDB:
         long[] words = new long[2];
         words[1] = 1L << 3;
         wordsField.set(bs, words);
-        Object _cf_result5 = bs.get(64 + 3);
+        boolean _cf_result5 = bs.get(64 + 3);
     }
 
     @Test
     public void testGet_LargeIndexDoesNotThrow_ReturnsFalse() {
-        Object _cf_result6 = instance.get(Integer.MAX_VALUE);
+        boolean _cf_result6 = instance.get(Integer.MAX_VALUE);
     }
 
     @Test
     public void testGet_BitBoundaryWordEdge63_ReturnsFalse() {
-        Object _cf_result7 = instance.get(63);
+        boolean _cf_result7 = instance.get(63);
     }
 
     @Test
     public void testGet_BitBoundaryWordEdge64_ReturnsFalse() {
-        Object _cf_result8 = instance.get(64);
+        boolean _cf_result8 = instance.get(64);
     }
 
     @Test
@@ -1109,7 +1109,7 @@ class TestBitSetLikeQuestDB:
         int nBits = 1_000_000;
         BitSet big = new BitSet(nBits);
         int last = nBits - 1;
-        Object _cf_result9 = big.get(last);
+        boolean _cf_result9 = big.get(last);
     }
 }
 """
@@ -1240,15 +1240,15 @@ def test_counters_assigned_in_source_order(self):
 public class FibTest {
     @Test
     void testA() {
-        Object _cf_result1 = Fibonacci.fibonacci(0);
+        int _cf_result1 = Fibonacci.fibonacci(0);
     }
     @Test
     void testB() {
-        Object _cf_result2 = Fibonacci.fibonacci(10);
+        int _cf_result2 = Fibonacci.fibonacci(10);
     }
     @Test
     void testC() {
-        Object _cf_result3 = Fibonacci.fibonacci(1);
+        int _cf_result3 = Fibonacci.fibonacci(1);
     }
 }
 """
@@ -1329,7 +1329,7 @@ def test_non_nested_assertions_all_replaced(self):
 public class FibTest {
     @Test
     void test() {
-        Object _cf_result1 = Fibonacci.fibonacci(0);
+        int _cf_result1 = Fibonacci.fibonacci(0);
     }
 }
 """
@@ -1362,11 +1362,11 @@ def test_reverse_replacement_preserves_all_positions(self):
 public class CalcTest {
     @Test
     void test() {
-        Object _cf_result1 = engine.compute(1);
-        Object _cf_result2 = engine.compute(2);
-        Object _cf_result3 = engine.compute(3);
-        Object _cf_result4 = engine.compute(4);
-        Object _cf_result5 = engine.compute(5);
+        int _cf_result1 = engine.compute(1);
+        int _cf_result2 = engine.compute(2);
+        int _cf_result3 = engine.compute(3);
+        int _cf_result4 = engine.compute(4);
+        int _cf_result5 = engine.compute(5);
     }
 }
 """
@@ -1400,8 +1400,8 @@ def test_mixed_assertions_all_removed(self):
 public class FibTest {
     @Test
     void test() {
-        Object _cf_result1 = Fibonacci.fibonacci(0);
-        Object _cf_result2 = Fibonacci.fibonacci(1);
+        int _cf_result1 = Fibonacci.fibonacci(0);
+        int _cf_result2 = Fibonacci.fibonacci(1);
     }
 }
 """
@@ -1461,7 +1461,7 @@ def test_single_assertion_exact_output(self):
 public class FibTest {
     @Test
     void test() {
-        Object _cf_result1 = Fibonacci.fibonacci(10);
+        int _cf_result1 = Fibonacci.fibonacci(10);
     }
 }
 """
@@ -1489,8 +1489,8 @@ def test_multiple_assertions_exact_output(self):
 public class CalcTest {
     @Test
     void test() {
-        Object _cf_result1 = calc.add(1, 2);
-        Object _cf_result2 = calc.add(3, 4);
+        int _cf_result1 = calc.add(1, 2);
+        int _cf_result2 = calc.add(3, 4);
     }
 }
 """
@@ -1546,12 +1546,12 @@ def test_invocation_counter_increments(self):
 public class FibTest {
     @Test
     void test1() {
-        Object _cf_result1 = Fibonacci.fibonacci(0);
+        int _cf_result1 = Fibonacci.fibonacci(0);
     }
 
     @Test
     void test2() {
-        Object _cf_result2 = Fibonacci.fibonacci(10);
+        int _cf_result2 = Fibonacci.fibonacci(10);
     }
 }
 """
@@ -1784,9 +1784,9 @@ class TestAllAssertionsRemoved:
 
     @Test
     void testFibonacci() {
-        Object _cf_result1 = Fibonacci.fibonacci(0);
-        Object _cf_result2 = Fibonacci.fibonacci(1);
-        Object _cf_result3 = Fibonacci.fibonacci(5);
+        int _cf_result1 = Fibonacci.fibonacci(0);
+        int _cf_result2 = Fibonacci.fibonacci(1);
+        int _cf_result3 = Fibonacci.fibonacci(5);
     }
 
     @Test
@@ -1846,7 +1846,7 @@ def test_preserves_non_assertion_code(self):
     void testAdd() {
         Calculator calc = new Calculator();
         int result = calc.setup();
-        Object _cf_result1 = calc.add(2, 3);
+        int _cf_result1 = calc.add(2, 3);
     }
 }
 """
@@ -1902,7 +1902,7 @@ def test_mixed_frameworks_all_removed(self):
 public class MixedTest {
     @Test
     void test() {
-        Object _cf_result1 = obj.target(1);
+        int _cf_result1 = obj.target(1);
     }
 }
 """

From 342a9c5b7474f35ee66d4676135711b0af57ace3 Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Wed, 25 Feb 2026 20:19:18 +0000
Subject: [PATCH 05/11] fix: increase Maven verify timeout from 120s to 300s
 for coverage runs

Multi-module projects like Guava require more time for the Maven verify
phase which runs compilation + instrumentation + test execution. The
120s minimum was causing timeouts for large projects.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 codeflash/languages/java/test_runner.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/codeflash/languages/java/test_runner.py b/codeflash/languages/java/test_runner.py
index 8dcaf401a..5c49ca075 100644
--- a/codeflash/languages/java/test_runner.py
+++ b/codeflash/languages/java/test_runner.py
@@ -427,8 +427,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:

From 171da454c7a16162daead1041a4cad4dc10aaa31 Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Wed, 25 Feb 2026 20:29:27 +0000
Subject: [PATCH 06/11] Optimize JavaAssertTransformer._infer_return_type
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Runtime improvement (primary): the optimized version cuts the measured wall-clock time from ~11.9 ms to ~5.23 ms (≈127% speedup). Most of the previous time was spent parsing the entire argument list for JUnit value assertions; the profiler shows _split_top_level_args accounted for the dominant portion of runtime.

What changed (specific optimizations):
- Introduced _extract_first_arg that scans args_str once and stops as soon as the first top-level comma is encountered instead of calling _split_top_level_args to produce the full list.
- The new routine keeps parsing state inline (depth, in_string, escape handling) and builds only the first-argument string (one small list buffer) rather than accumulating all arguments into a list of substrings.
- Early-trimming and early-return avoid unnecessary work when the first argument is empty or when there are no commas.

Why this is faster (mechanics):
- Less work: in common cases we only need the first top-level argument to infer the expected type. Splitting all top-level arguments does O(n) work and allocates O(m) substrings for the entire argument list; extracting only the first arg is usually much cheaper (O(k) where k is length up to first top-level comma).
- Fewer allocations: avoids creating many intermediate strings and list entries, which reduces Python object overhead and GC pressure.
- Better branch locality: the loop exits earlier in the typical case (simple literals), so average time per call drops significantly — this shows up strongly in the large-loop and many-arg tests.

Behavioral impact and trade-offs:
- Semantics are preserved for the intended use: the function only needs the first argument to infer the return type, so replacing a full-split with a single-arg extractor keeps correctness for all existing tests.
- Microbenchmarks for very trivial cases (e.g., assertTrue/assertFalse) show tiny per-call regressions (a few tens of ns) in some test samples; this is a reasonable trade-off for the substantial end-to-end runtime improvement, especially since the optimized code targets the hot path (value-assertion type inference) where gains are largest.

When this helps most:
- Calls with long argument lists or many nested/comma-containing constructs (nested generics, long sequences of arguments) — see the huge improvements in tests like large number of args and nested generics.
- Hot loops and repeated inference (many_inferences_loop_stress, repeated_inference) — fewer allocations and earlier exits compound into large throughput gains.

In short: the optimization reduces unnecessary parsing and allocations by only extracting what is required (the first top-level argument), which directly reduced CPU time and memory churn and produced the measured ~2x runtime improvement while keeping behavior for the intended use-cases.
---
 codeflash/languages/java/remove_asserts.py | 60 ++++++++++++++++++++--
 1 file changed, 56 insertions(+), 4 deletions(-)

diff --git a/codeflash/languages/java/remove_asserts.py b/codeflash/languages/java/remove_asserts.py
index d89e76cce..2e0d06864 100644
--- a/codeflash/languages/java/remove_asserts.py
+++ b/codeflash/languages/java/remove_asserts.py
@@ -941,15 +941,15 @@ def _infer_type_from_assertion_args(self, original_text: str, method: str) -> st
         elif args_str.endswith(")"):
             args_str = args_str[:-1]
 
-        # Split top-level args (respecting parens, strings, generics)
-        args = self._split_top_level_args(args_str)
-        if not args:
+        # 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"
 
         # assertEquals has (expected, actual) or (expected, actual, message/delta)
         # Some overloads have (message, expected, actual) in JUnit 4 but JUnit 5 uses (expected, actual[, message])
         # Try the first argument as the expected value
-        expected = args[0].strip()
+        expected = first_arg.strip()
 
         return self._type_from_literal(expected)
 
@@ -1108,6 +1108,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.

From 52ec0a904c4d4a2f2e345c4fb33412127cdc88ae Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Wed, 25 Feb 2026 20:34:12 +0000
Subject: [PATCH 07/11] Optimize
 JavaAssertTransformer._infer_type_from_assertion_args
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Runtime improvement (primary): the optimized version runs ~11% faster overall (10.3ms -> 9.23ms). Line-profiles show the hot work (argument splitting and literal checks) is measurably reduced.

What changed (concrete):
- Added a fast-path in _split_top_level_args: if the args string contains none of the "special" delimiters (quotes, braces, parens), we skip the character-by-character parser and return either args_str.split(",") or [args_str].
- Moved several literal/cast regexes into __init__ as precompiled attributes (self._FLOAT_LITERAL_RE, self._DOUBLE_LITERAL_RE, self._LONG_LITERAL_RE, self._INT_LITERAL_RE, self._CHAR_LITERAL_RE, self._cast_re) and replaced re.match(...) for casts with self._cast_re.match(...).

Why this speeds things up:
- str.split is implemented in C and is orders of magnitude faster than a Python-level loop that iterates characters, manages stack depth, and joins fragments. The fast-path catches the common simple cases (no nested parentheses/quotes/generics) and lets the interpreter use the highly-optimized C split, which is why very large comma-separated inputs show the biggest wins (e.g., the 1000-arg test goes from ~1.39ms to ~67.5μs).
- Precompiling regexes removes repeated compilation overhead and lets .match be executed directly on a compiled object. The original code used re.match(...) in-place for cast detection which implicitly compiles the pattern or goes through the module-level cache; using a stored compiled pattern is cheaper and eliminates that runtime cost.
- Combined, these changes reduce the time spent inside _split_top_level_args and _type_from_literal (the line profilers show reduced wall time for those functions), producing the measured global runtime improvement.

Behavioral/compatibility notes:
- The fast-path preserves original behavior: when no special delimiter is present it simply splits on commas (or returns a single entry), otherwise it falls back to the full, safe parser that respects nested delimiters and strings.
- Some microbenchmarks regress slightly (a few single-case timings in the annotated tests are a bit slower); this is expected because we add a small _special_re.search check for every call. The overall trade-off was accepted because it yields substantial savings in the common and expensive cases (especially large/simple comma-separated argument lists).
- The optimization is most valuable when this function is exercised many times or on long/simple argument lists (hot paths that produce many simple comma-separated tokens). It is neutral or slightly negative for a handful of small or highly-nested inputs, but those are rare in the benchmarks.

Tests and workload guidance:
- Big wins: large-scale, many-argument inputs or many repeated calls where arguments are simple comma-separated literals (annotated tests show up to ~20x speedups for such cases).
- No/low impact: complex first arguments with nested parentheses/generics or many quoted strings — the safe parser still runs there, so correctness is preserved; timings remain similar.
- Small regressions: a few microbench cases (very short inputs or certain char-literal checks) are marginally slower due to the extra quick search, but these regressions are small relative to the global runtime improvement.

Summary:
By routing simple/common inputs to str.split (C-level speed) and eliminating per-call regex compilation for literal/cast detection, the optimized code reduces time in the hot parsing and literal-detection paths, producing the observed ~11% runtime improvement while maintaining correctness for nested/quoted input via the fallback parser.
---
 codeflash/languages/java/remove_asserts.py | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/codeflash/languages/java/remove_asserts.py b/codeflash/languages/java/remove_asserts.py
index d89e76cce..ee8c4d713 100644
--- a/codeflash/languages/java/remove_asserts.py
+++ b/codeflash/languages/java/remove_asserts.py
@@ -198,6 +198,15 @@ 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.
 
@@ -972,13 +981,22 @@ def _type_from_literal(self, value: str) -> str:
         if value.startswith('"'):
             return "String"
         # Cast expression like (byte)0, (short)1
-        cast_match = re.match(r"^\((\w+)\)", value)
+        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] = []

From 15a37fa0ff41051edd5894d02a4c18ad94d8009f Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Wed, 25 Feb 2026 20:37:13 +0000
Subject: [PATCH 08/11] test: update instrumentation test for assertion type
 inference

The behavior mode instrumentation test expected `Object _cf_result1`
but after the type inference fix, assertEquals(4, call()) now produces
`int _cf_result1 = (int)_cf_result1_1`.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 tests/test_languages/test_java/test_instrumentation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_languages/test_java/test_instrumentation.py b/tests/test_languages/test_java/test_instrumentation.py
index 34d6da567..7dc891330 100644
--- a/tests/test_languages/test_java/test_instrumentation.py
+++ b/tests/test_languages/test_java/test_instrumentation.py
@@ -1324,7 +1324,7 @@ def test_instrument_generated_test_behavior_mode(self):
                 }
             }
         }
-        Object _cf_result1 = _cf_result1_1;
+        int _cf_result1 = (int)_cf_result1_1;
     }
 }
 """

From 23d12fbfab6d29964eca4c451f72b53a6a6c0fb1 Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Wed, 25 Feb 2026 20:52:23 +0000
Subject: [PATCH 09/11] Optimize JavaAssertTransformer._generate_replacement
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Primary benefit: the optimized version reduces runtime from ~903μs to ~802μs — a 12% speedup — by lowering per-iteration and attribute-access overhead in the hot path of _generate_replacement.

What changed (concrete, low-level):
- Cached attribute lookups into locals:
  - self.invocation_counter → local inv, written back once at the end.
  - assertion.leading_whitespace and assertion.target_calls → leading_ws and calls locals.
  Caching avoids repeated attribute reads/writes which are relatively expensive in Python.
- Removed a per-iteration branch by handling the first target call separately:
  - The original loop used if i == 0 every iteration (via enumerate). The optimized code emits the first line once, then loops the remaining calls without a conditional.
  This eliminates an O(n) conditional check across many iterations.
- Reduced formatting overhead for loop-generated variable names:
  - var_name is built with "_cf_result" + str(inv) instead of using an f-string inside the loop (fewer formatting operations).
- Minor local micro-optimizations in _infer_return_type:
  - Replaced the small "in (a, b)" checks with equivalent chained comparisons (method == "x" or method == "y"), reducing tuple creation/containment checks.
- Exception-replacement counter handling: moved to a local increment-and-write-back pattern (same semantics, fewer attribute writes).

Why this speeds things up:
- Attribute access and writes (self.foo / assertion.attr) cost significantly more than local variable access. By doing those once per call and using locals inside tight loops we reduce Python bytecode operations dramatically.
- Removing the per-iteration i == 0 branch eliminates one conditional per target call; for large lists this reduces branching overhead and improves instruction cache behavior.
- Minimizing string formatting and concatenation inside a hot loop reduces temporary allocations and joins fewer intermediate values.
- The profiler and tests show the biggest gains appear when there are many target_calls (1000-call test: ~240μs → ~202μs, ~19% faster), matching these optimizations’ effect on O(n) behavior.

Behavioral impact and correctness:
- The observable behavior (variable names, formatting, invocation_counter progression, and exception-handling output) is preserved. The counter is still incremented the same number of times and persists across calls.
- Exception handling logic is unchanged semantically; only the internal counter updates were made more efficient.

Trade-offs (noted regressions and why they’re acceptable):
- A few small test cases show tiny slowdowns (single very-small assertions, some assertDoesNotThrow paths). These are microsecond-level regressions (often <0.1–0.2μs) and are an acceptable trade-off for sizable improvements in the common hot path (large lists and repeated invocations).
- The optimizations prioritize reducing per-iteration overhead; therefore workloads dominated by many target calls or repeated invocations benefit most. Small or one-off assertions will see negligible change.

Where this helps most (based on tests and profiler):
- Hot paths that iterate many times over assertion.target_calls (large test files or transformations producing hundreds/thousands of captures).
- Repeated uses of the same transformer instance where invocation_counter accumulates across many calls.
- The annotated tests and line profiler confirm the speedup is concentrated in _generate_replacement’s loop and that the large-list tests (n=1000) get the biggest absolute and relative improvement.

In short: the optimized code reduces attribute and branching overhead in the hot loop, cutting allocation/bytecode work per target call — which yields the observed 12% runtime improvement and up to ~19% on large inputs while preserving behavior.
---
 codeflash/languages/java/remove_asserts.py | 40 +++++++++++++++-------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/codeflash/languages/java/remove_asserts.py b/codeflash/languages/java/remove_asserts.py
index d89e76cce..4c65d96ee 100644
--- a/codeflash/languages/java/remove_asserts.py
+++ b/codeflash/languages/java/remove_asserts.py
@@ -904,11 +904,11 @@ def _infer_return_type(self, assertion: AssertionMatch) -> str:
         method = assertion.assertion_method
 
         # assertTrue/assertFalse always deal with boolean values
-        if method in ("assertTrue", "assertFalse"):
+        if method == "assertTrue" or method == "assertFalse":
             return "boolean"
 
         # assertNull/assertNotNull — keep Object (reference type)
-        if method in ("assertNull", "assertNotNull"):
+        if method == "assertNull" or method == "assertNotNull":
             return "Object"
 
         # For assertEquals/assertNotEquals/assertSame, try to infer from the expected literal
@@ -1039,18 +1039,32 @@ def _generate_replacement(self, assertion: AssertionMatch) -> str:
         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}{return_type} {var_name} = {call.full_call};")
-            else:
+        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)
 
     def _generate_exception_replacement(self, assertion: AssertionMatch) -> str:
@@ -1068,8 +1082,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")
 

From cd406b9514e4ae8fff3b21073fee1a8cea3d5fb7 Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Wed, 25 Feb 2026 22:07:34 +0000
Subject: [PATCH 10/11] refactor: remove CLI-side Java import band-aid

Remove _COMMON_JAVA_IMPORTS, ensure_common_java_imports(), and _add_import()
from instrumentation.py. The root cause is now fixed in the AI service
(codeflash-internal#2443) which adds comprehensive stdlib import postprocessing
before tree-sitter validation in the testgen pipeline.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 codeflash/languages/java/instrumentation.py | 59 ---------------------
 1 file changed, 59 deletions(-)

diff --git a/codeflash/languages/java/instrumentation.py b/codeflash/languages/java/instrumentation.py
index 81e125ae6..5bdfb0b14 100644
--- a/codeflash/languages/java/instrumentation.py
+++ b/codeflash/languages/java/instrumentation.py
@@ -1260,35 +1260,6 @@ def remove_instrumentation(source: str) -> str:
     return source
 
 
-_COMMON_JAVA_IMPORTS = {
-    "Arrays": "import java.util.Arrays;",
-    "List": "import java.util.List;",
-    "ArrayList": "import java.util.ArrayList;",
-    "Map": "import java.util.Map;",
-    "HashMap": "import java.util.HashMap;",
-    "Set": "import java.util.Set;",
-    "HashSet": "import java.util.HashSet;",
-    "Collections": "import java.util.Collections;",
-    "Collectors": "import java.util.stream.Collectors;",
-    "Random": "import java.util.Random;",
-    "BigDecimal": "import java.math.BigDecimal;",
-    "BigInteger": "import java.math.BigInteger;",
-}
-
-
-def ensure_common_java_imports(test_code: str) -> str:
-    for class_name, import_stmt in _COMMON_JAVA_IMPORTS.items():
-        if not re.search(rf"\b{class_name}\b", test_code):
-            continue
-        if import_stmt in test_code:
-            continue
-        package = import_stmt.split()[1].rsplit(".", 1)[0]
-        if f"import {package}.*;" in test_code:
-            continue
-        test_code = _add_import(test_code, import_stmt)
-    return test_code
-
-
 def instrument_generated_java_test(
     test_code: str,
     function_name: str,
@@ -1319,7 +1290,6 @@ def instrument_generated_java_test(
     from codeflash.languages.java.remove_asserts import transform_java_assertions
 
     test_code = transform_java_assertions(test_code, function_name, qualified_name)
-    test_code = ensure_common_java_imports(test_code)
 
     # Extract class name from the test code
     # Use pattern that starts at beginning of line to avoid matching words in comments
@@ -1355,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)

From dd3bdaf809b1756bf8e9e106be3e17b4bb7b0cac Mon Sep 17 00:00:00 2001
From: Mohamed Ashraf <ashraf@codeflash.ai>
Date: Thu, 26 Feb 2026 00:45:53 +0000
Subject: [PATCH 11/11] fix: handle non-standard JUnit parameterized test name
 patterns in loop_index parsing

The previous code assumed test names with brackets always follow the
pattern "testName[ N ]" (space after bracket). JUnit 5 parameterized
tests produce names like "testName(int)[1]" or "testName(String)[label]"
which caused a ValueError crash when parsing the loop index.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 codeflash/verification/parse_test_output.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/codeflash/verification/parse_test_output.py b/codeflash/verification/parse_test_output.py
index deb7d3a4b..b4fda1e2d 100644
--- a/codeflash/verification/parse_test_output.py
+++ b/codeflash/verification/parse_test_output.py
@@ -805,7 +805,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: