fix: handle assertThrows variable assignment in Java instrumentation

When assertThrows was assigned to a variable to validate exception
properties, the transformation generated invalid Java syntax by
replacing the assertThrows call with try-catch while leaving the
variable assignment intact.

Example of invalid output:
  IllegalArgumentException e = try { code(); } catch (Exception) {}

This fix detects variable assignments, extracts the exception type
from assertThrows arguments, and generates proper exception capture:
  IllegalArgumentException e = null;
  try { code(); } catch (IllegalArgumentException _cf_caught1) { e = _cf_caught1; } catch (Exception _cf_ignored1) {}

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: mohamedashrraf222

codeflash/languages/java/remove_asserts.py

@@ -166,6 +166,9 @@ class AssertionMatch:
     original_text: str = ""
     is_exception_assertion: bool = False
     lambda_body: str | None = None  # For assertThrows lambda content
+    variable_type: str | None = None  # Type of assigned variable (e.g., "IllegalArgumentException")
+    variable_name: str | None = None  # Name of assigned variable (e.g., "exception")
+    exception_class: str | None = None  # Exception class from assertThrows args
 
 
 class JavaAssertTransformer:
@@ -326,12 +329,32 @@ def _find_junit_assertions(self, source: str) -> list[AssertionMatch]:
             target_calls = self._extract_target_calls(args_content, match.end())
             is_exception = assertion_method in JUNIT5_EXCEPTION_ASSERTIONS
 
-            # For assertThrows, extract the lambda body
+            # For assertThrows, extract the lambda body and exception class
             lambda_body = None
+            exception_class = None
             if is_exception and assertion_method == "assertThrows":
                 lambda_body = self._extract_lambda_body(args_content)
+                exception_class = self._extract_exception_class(args_content)
+
+            # Check if assertion is assigned to a variable
+            var_type, var_name = self._detect_variable_assignment(source, start_pos)
+
+            # If variable assignment detected, adjust start_pos to include the entire line
+            actual_start = start_pos
+            actual_leading_ws = leading_ws
+            if var_type:
+                # Find the start of the line (beginning of variable declaration)
+                line_start = source.rfind("\n", 0, start_pos)
+                if line_start == -1:
+                    line_start = 0
+                else:
+                    line_start += 1
+                actual_start = line_start
+                # Extract the actual leading whitespace from the start of the line
+                line_content = source[line_start:start_pos]
+                actual_leading_ws = line_content[:len(line_content) - len(line_content.lstrip())]
 
-            original_text = source[start_pos:end_pos]
+            original_text = source[actual_start:end_pos]
 
             # Determine statement type based on detected framework
             detected = self._detected_framework or "junit5"
@@ -342,15 +365,18 @@ def _find_junit_assertions(self, source: str) -> list[AssertionMatch]:
 
             assertions.append(
                 AssertionMatch(
-                    start_pos=start_pos,
+                    start_pos=actual_start,
                     end_pos=end_pos,
                     statement_type=stmt_type,
                     assertion_method=assertion_method,
                     target_calls=target_calls,
-                    leading_whitespace=leading_ws,
+                    leading_whitespace=actual_leading_ws,
                     original_text=original_text,
                     is_exception_assertion=is_exception,
                     lambda_body=lambda_body,
+                    variable_type=var_type,
+                    variable_name=var_name,
+                    exception_class=exception_class,
                 )
             )
 
@@ -580,6 +606,85 @@ def _extract_target_calls(self, content: str, base_offset: int) -> list[TargetCa
 
         return target_calls
 
+    def _detect_variable_assignment(self, source: str, assertion_start: int) -> tuple[str | None, str | None]:
+        """Check if assertion is assigned to a variable.
+
+        Detects patterns like:
+            IllegalArgumentException exception = assertThrows(...)
+            Exception ex = assertThrows(...)
+
+        Args:
+            source: The full source code.
+            assertion_start: Start position of the assertion.
+
+        Returns:
+            Tuple of (variable_type, variable_name) or (None, None).
+
+        """
+        # Look backwards from assertion_start to beginning of line
+        line_start = source.rfind("\n", 0, assertion_start)
+        if line_start == -1:
+            line_start = 0
+        else:
+            line_start += 1
+
+        line_before_assert = source[line_start:assertion_start]
+
+        # Pattern: Type varName = assertXxx(...)
+        # Handle generic types: Type<Generic> varName = ...
+        pattern = r"(\w+(?:<[^>]+>)?)\s+(\w+)\s*=\s*$"
+        match = re.search(pattern, line_before_assert)
+
+        if match:
+            var_type = match.group(1).strip()
+            var_name = match.group(2).strip()
+            return var_type, var_name
+
+        return None, None
+
+    def _extract_exception_class(self, args_content: str) -> str | None:
+        """Extract exception class from assertThrows arguments.
+
+        Args:
+            args_content: Content inside assertThrows parentheses.
+
+        Returns:
+            Exception class name (e.g., "IllegalArgumentException") or None.
+
+        Example:
+            assertThrows(IllegalArgumentException.class, ...) -> "IllegalArgumentException"
+
+        """
+        # First argument is the exception class reference (e.g., "IllegalArgumentException.class")
+        # Split by comma, but respect nested parentheses and generics
+        depth = 0
+        current = []
+        parts = []
+
+        for char in args_content:
+            if char in "(<":
+                depth += 1
+                current.append(char)
+            elif char in ")>":
+                depth -= 1
+                current.append(char)
+            elif char == "," and depth == 0:
+                parts.append("".join(current).strip())
+                current = []
+            else:
+                current.append(char)
+
+        if current:
+            parts.append("".join(current).strip())
+
+        if parts:
+            exception_arg = parts[0].strip()
+            # Remove .class suffix
+            if exception_arg.endswith(".class"):
+                return exception_arg[:-6].strip()
+
+        return None
+
     def _extract_lambda_body(self, content: str) -> str | None:
         """Extract the body of a lambda expression from assertThrows arguments.
 
@@ -745,29 +850,53 @@ def _generate_exception_replacement(self, assertion: AssertionMatch) -> str:
         To:
             try { calculator.divide(1, 0); } catch (Exception _cf_ignored1) {}
 
+        For variable assignments:
+            IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> code());
+        To:
+            IllegalArgumentException ex = null;
+            try { code(); } catch (IllegalArgumentException e) { ex = e; } catch (Exception _cf_ignored1) {}
+
         """
         self.invocation_counter += 1
 
+        # Extract code to run from lambda body or target calls
+        code_to_run = None
         if assertion.lambda_body:
-            # Extract the actual code from the lambda
             code_to_run = assertion.lambda_body
             if not code_to_run.endswith(";"):
                 code_to_run += ";"
-            return (
-                f"{assertion.leading_whitespace}try {{ {code_to_run} }} "
-                f"catch (Exception _cf_ignored{self.invocation_counter}) {{}}"
-            )
-
-        # If no lambda body found, try to extract from target calls
-        if assertion.target_calls:
+        elif assertion.target_calls:
             call = assertion.target_calls[0]
+            code_to_run = call.full_call + ";"
+
+        if not code_to_run:
+            # Fallback: comment out the assertion
+            return f"{assertion.leading_whitespace}// Removed assertThrows: could not extract callable"
+
+        # Check if assertion is assigned to a variable
+        if assertion.variable_name and assertion.variable_type:
+            # Generate proper exception capture with variable assignment
+            exception_type = assertion.exception_class or assertion.variable_type
+            var_name = assertion.variable_name
+
+            # Use a unique catch variable name to avoid conflicts
+            catch_var = f"_cf_caught{self.invocation_counter}"
+
+            # Get base indentation from leading whitespace (without newlines)
+            base_indent = assertion.leading_whitespace.lstrip("\n\r")
+
             return (
-                f"{assertion.leading_whitespace}try {{ {call.full_call}; }} "
+                f"{assertion.leading_whitespace}{assertion.variable_type} {var_name} = null;\n"
+                f"{base_indent}try {{ {code_to_run} }} "
+                f"catch ({exception_type} {catch_var}) {{ {var_name} = {catch_var}; }} "
                 f"catch (Exception _cf_ignored{self.invocation_counter}) {{}}"
             )
 
-        # Fallback: comment out the assertion
-        return f"{assertion.leading_whitespace}// Removed assertThrows: could not extract callable"
+        # No variable assignment, use simple try-catch
+        return (
+            f"{assertion.leading_whitespace}try {{ {code_to_run} }} "
+            f"catch (Exception _cf_ignored{self.invocation_counter}) {{}}"
+        )
 
 
 def transform_java_assertions(source: str, function_name: str, qualified_name: str | None = None) -> str:

tests/test_java_assertion_removal.py

@@ -1255,3 +1255,106 @@ def test_concurrent_assertion_with_assertj(self):
 }"""
         result = transform_java_assertions(source, "incrementAndGet")
         assert result == expected
+
+
+class TestAssertThrowsVariableAssignment:
+    """Tests for assertThrows with variable assignment (Issue: exception handling instrumentation bug)."""
+
+    def test_assert_throws_with_variable_assignment_expression_lambda(self):
+        """Test assertThrows assigned to variable with expression lambda."""
+        source = """\
+@Test
+void testNegativeInput() {
+    IllegalArgumentException exception = assertThrows(
+        IllegalArgumentException.class,
+        () -> calculator.fibonacci(-1)
+    );
+    assertEquals("Negative input not allowed", exception.getMessage());
+}"""
+        expected = """\
+@Test
+void testNegativeInput() {
+    IllegalArgumentException exception = null;
+    try { calculator.fibonacci(-1); } catch (IllegalArgumentException _cf_caught1) { exception = _cf_caught1; } catch (Exception _cf_ignored1) {}
+    assertEquals("Negative input not allowed", exception.getMessage());
+}"""
+        result = transform_java_assertions(source, "fibonacci")
+        assert result == expected
+
+    def test_assert_throws_with_variable_assignment_block_lambda(self):
+        """Test assertThrows assigned to variable with block lambda."""
+        source = """\
+@Test
+void testInvalidOperation() {
+    ArithmeticException ex = assertThrows(ArithmeticException.class, () -> {
+        calculator.divide(10, 0);
+    });
+    assertEquals("Division by zero", ex.getMessage());
+}"""
+        expected = """\
+@Test
+void testInvalidOperation() {
+    ArithmeticException ex = null;
+    try { calculator.divide(10, 0); } catch (ArithmeticException _cf_caught1) { ex = _cf_caught1; } catch (Exception _cf_ignored1) {}
+    assertEquals("Division by zero", ex.getMessage());
+}"""
+        result = transform_java_assertions(source, "divide")
+        assert result == expected
+
+    def test_assert_throws_with_variable_assignment_generic_exception(self):
+        """Test assertThrows with generic Exception type."""
+        source = """\
+@Test
+void testGenericException() {
+    Exception e = assertThrows(Exception.class, () -> processor.process(null));
+    assertNotNull(e.getMessage());
+}"""
+        expected = """\
+@Test
+void testGenericException() {
+    Exception e = null;
+    try { processor.process(null); } catch (Exception _cf_caught1) { e = _cf_caught1; } catch (Exception _cf_ignored1) {}
+    assertNotNull(e.getMessage());
+}"""
+        result = transform_java_assertions(source, "process")
+        assert result == expected
+
+    def test_assert_throws_without_variable_assignment(self):
+        """Test assertThrows without variable assignment still works (no regression)."""
+        source = """\
+@Test
+void testThrowsException() {
+    assertThrows(IllegalArgumentException.class, () -> calculator.fibonacci(-1));
+}"""
+        expected = """\
+@Test
+void testThrowsException() {
+    try { calculator.fibonacci(-1); } catch (Exception _cf_ignored1) {}
+}"""
+        result = transform_java_assertions(source, "fibonacci")
+        assert result == expected
+
+    def test_assert_throws_with_variable_and_multi_line_lambda(self):
+        """Test assertThrows with variable assignment and multi-line lambda."""
+        source = """\
+@Test
+void testComplexException() {
+    IllegalStateException exception = assertThrows(
+        IllegalStateException.class,
+        () -> {
+            processor.initialize();
+            processor.execute();
+        }
+    );
+    assertTrue(exception.getMessage().contains("not initialized"));
+}"""
+        expected = """\
+@Test
+void testComplexException() {
+    IllegalStateException exception = null;
+    try { processor.initialize();
+            processor.execute(); } catch (IllegalStateException _cf_caught1) { exception = _cf_caught1; } catch (Exception _cf_ignored1) {}
+    assertTrue(exception.getMessage().contains("not initialized"));
+}"""
+        result = transform_java_assertions(source, "execute")
+        assert result == expected