⚡️ Speed up method JavaAssertTransformer._find_junit_assertions by 22% in PR #1295 (feat/java-remove-asserts-transformer)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1295

If you approve this dependent PR, these changes will be merged into the original PR branch feat/java-remove-asserts-transformer.

This PR will be automatically closed if the original PR is merged.

---

📄 22% (0.22x) speedup for JavaAssertTransformer._find_junit_assertions in codeflash/languages/java/remove_asserts.py

⏱️ Runtime : 9.88 milliseconds → 8.08 milliseconds (best of 100 runs)

📝 Explanation and details

The optimized code achieves a 22% runtime improvement by eliminating repeated regex pattern compilation overhead in hot paths.

Key Optimizations:

1. Pre-compiled Regex Patterns (Primary Speedup):

   • Moved three re.compile() calls from method bodies to __init__:
     • _junit_pattern: Compiled once instead of on every _find_junit_assertions() call
     • _target_call_pattern: Compiled once instead of on every _extract_target_calls() call
     • _lambda_pattern: Compiled once instead of on every _extract_lambda_body() call
   • The line profiler shows _find_junit_assertions() dropped from 78.8ms to 63.8ms (19% faster), with regex compilation time eliminated from the 4.2% hotspot
   • Similarly, _extract_target_calls() improved from 18.6ms to 12.1ms (35% faster), removing its 34.8% regex compilation overhead

2. Optimized String Indexing in _find_balanced_parens():

   • Caches len(code) as code_len to avoid repeated function calls in the tight loop
   • Restructured escape sequence checking to avoid redundant code[pos - 1] lookups
   • Reduced per-character overhead in the parser, improving from 26.7ms to 23.4ms (12% faster)

3. Simplified Lambda Body Extraction:

   • Replaced nested content[body_start:].index("{") calls with single content.index("{", body_start)
   • Reduced method time from 456μs to 144μs (68% faster on assertThrows cases)

Why This Works:

• Regex compilation in Python is expensive (involves pattern parsing, DFA construction). By compiling patterns once during initialization rather than on every method call, we eliminate this overhead from the critical path
• The test results show consistent 30-70% improvements across all test cases, with the largest gains on simpler assertions where regex compilation dominated runtime
• These optimizations are particularly effective for code analysis tools that process many assertions repeatedly (see the test_large_scale_many_assertions_under_limit improving from 1.54ms to 1.13ms)

The changes preserve all functionality while significantly improving performance for Java assertion analysis workflows.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 206 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	92.9%

🌀 Click to see Generated Regression Tests

from typing import List

# imports
import pytest  # used for our unit tests
from codeflash.languages.java.remove_asserts import JavaAssertTransformer

def test_basic_assert_equals_detects_target_call():
    # Ensure a simple assertEquals with a target() call is detected and parsed.
    transformer = JavaAssertTransformer(function_name="target")  # instance method requires real instance

    # Source has leading whitespace before the assertion to validate leading_whitespace capture.
    source = '   assertEquals(42, target(arg1, "str"));'
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 29.0μs -> 20.7μs (40.2% faster)

    match = results[0]

    target_call = match.target_calls[0]

def test_qualified_assertions_and_receiver_detection():
    # Verify assertions with qualified names and receiver detection for method calls.
    transformer = JavaAssertTransformer(function_name="target")

    # Use qualified assertion (Assertions.assertTrue) and call target on an object receiver.
    source = "Assertions.assertTrue(obj.target());"
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.7μs -> 15.2μs (42.5% faster)
    match = results[0]
    tc = match.target_calls[0]

def test_assert_throws_extracts_lambda_expression_and_block_body():
    # Validate special handling for assertThrows with both expression and block lambdas.
    transformer = JavaAssertTransformer(function_name="target")

    # Expression-style lambda: () -> target()
    source_expr = "assertThrows(Exception.class, () -> target());"
    codeflash_output = transformer._find_junit_assertions(source_expr); res_expr = codeflash_output # 33.9μs -> 25.8μs (31.4% faster)
    m_expr = res_expr[0]

    # Block-style lambda: () -> { target(); }
    source_block = "assertThrows(Exception.class, () -> { target(); });"
    codeflash_output = transformer._find_junit_assertions(source_block); res_block = codeflash_output # 28.9μs -> 22.5μs (28.1% faster)
    m_block = res_block[0]

def test_incomplete_parentheses_are_skipped_and_do_not_crash():
    # An assertion with unbalanced parentheses should be ignored (returned list empty).
    transformer = JavaAssertTransformer(function_name="target")
    # Missing closing paren for target invocation -> will make balanced paren search fail
    source = "assertEquals(1, target("
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 10.3μs -> 6.65μs (55.4% faster)

def test_string_and_char_literals_with_parentheses_do_not_break_parsing():
    # Parentheses inside string or char literals should not confuse balanced-paren parsing.
    transformer = JavaAssertTransformer(function_name="target")

    # String literal contains parentheses and char literal contains a parenthesis char.
    source = 'assertEquals(")( and parentheses", target());\nassertEquals(\')\', target());'
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 40.7μs -> 31.1μs (31.2% faster)

def test_multiple_assertions_detected_with_correct_positions():
    # Ensure multiple assertions in a file are all detected and positions reflect their location.
    transformer = JavaAssertTransformer(function_name="target")

    # Two assertions on separate lines; positions should increase.
    source = "assertEquals(1, target());\n   Assert.assertEquals(2, target(2));\n"
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 33.4μs -> 24.0μs (39.1% faster)

    first, second = results[0], results[1]

def test_assert_with_object_and_class_receivers_and_full_call_text():
    # Validate detection when target is invoked on a class or instance receiver.
    transformer = JavaAssertTransformer(function_name="target")

    source = "Assert.assertEquals(msg, MyClass.target(x, y));\nAssertions.assertEquals(msg2, instance.target());"
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 40.0μs -> 30.0μs (33.5% faster)

    # First target call should have receiver 'MyClass'
    tc0 = results[0].target_calls[0]

    # Second target call should have receiver 'instance'
    tc1 = results[1].target_calls[0]

def test_large_scale_many_assertions_under_limit():
    # Stress-test with many assertions to ensure scalability (keeps under 1000 elements as required).
    transformer = JavaAssertTransformer(function_name="target")

    # Build 200 assertions (well under 1000). Using repetition to simulate a larger file.
    count = 200
    lines: List[str] = ["assertTrue(target());" for _ in range(count)]
    source = "\n".join(lines)

    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 1.54ms -> 1.13ms (37.2% faster)

    # Check a few random properties: each match should include a target call and be of assertion_method assertTrue.
    for m in results:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from codeflash.languages.java.remove_asserts import (JUNIT5_ALL_ASSERTIONS,
                                                     AssertionMatch,
                                                     JavaAssertTransformer,
                                                     TargetCall)

def test_find_junit_assertions_simple_assertEquals():
    """Test finding a simple assertEquals assertion with static import."""
    source = "assertEquals(expected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.6μs -> 14.7μs (46.5% faster)

def test_find_junit_assertions_qualified_assert_class():
    """Test finding assertions with Assert. prefix (JUnit 4 style)."""
    source = "Assert.assertEquals(expected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 20.7μs -> 14.3μs (45.0% faster)

def test_find_junit_assertions_qualified_assertions_class():
    """Test finding assertions with Assertions. prefix (JUnit 5 style)."""
    source = "Assertions.assertEquals(expected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 20.4μs -> 14.3μs (42.7% faster)

def test_find_junit_assertions_assertTrue():
    """Test finding assertTrue assertion."""
    source = "assertTrue(condition);"
    transformer = JavaAssertTransformer("condition")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 16.7μs -> 10.3μs (61.2% faster)

def test_find_junit_assertions_assertFalse():
    """Test finding assertFalse assertion."""
    source = "assertFalse(condition);"
    transformer = JavaAssertTransformer("condition")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 16.5μs -> 10.1μs (64.0% faster)

def test_find_junit_assertions_assertNull():
    """Test finding assertNull assertion."""
    source = "assertNull(object);"
    transformer = JavaAssertTransformer("object")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 15.8μs -> 9.24μs (70.8% faster)

def test_find_junit_assertions_assertNotNull():
    """Test finding assertNotNull assertion."""
    source = "assertNotNull(object);"
    transformer = JavaAssertTransformer("object")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 15.1μs -> 9.22μs (63.8% faster)

def test_find_junit_assertions_multiple_assertions():
    """Test finding multiple assertions in the same source."""
    source = """
        assertEquals(5, getValue());
        assertTrue(isValid());
        assertFalse(isEmpty());
    """
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 38.9μs -> 28.3μs (37.5% faster)

def test_find_junit_assertions_with_leading_whitespace():
    """Test that leading whitespace is captured correctly."""
    source = "        assertEquals(5, result);"
    transformer = JavaAssertTransformer("result")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 18.0μs -> 11.5μs (56.3% faster)

def test_find_junit_assertions_extracting_target_calls():
    """Test that target function calls are extracted from assertion arguments."""
    source = "assertEquals(5, getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.5μs -> 15.2μs (41.3% faster)

def test_find_junit_assertions_with_arguments():
    """Test extracting target calls with arguments."""
    source = "assertEquals(5, getValue(param1, param2));"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 28.3μs -> 20.1μs (40.7% faster)

def test_find_junit_assertions_assertThrows():
    """Test finding assertThrows exception assertion."""
    source = "assertThrows(IllegalArgumentException.class, () -> getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 41.4μs -> 31.9μs (29.7% faster)

def test_find_junit_assertions_assertDoesNotThrow():
    """Test finding assertDoesNotThrow assertion."""
    source = "assertDoesNotThrow(() -> getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 22.3μs -> 15.9μs (40.1% faster)

def test_find_junit_assertions_with_receiver_object():
    """Test extracting target calls with receiver object."""
    source = "assertEquals(5, obj.getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 22.9μs -> 16.2μs (41.1% faster)

def test_find_junit_assertions_empty_source():
    """Test with empty source code."""
    source = ""
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 4.82μs -> 1.60μs (201% faster)

def test_find_junit_assertions_no_assertions():
    """Test source with no assertions."""
    source = "int x = 5; String y = getValue(); return x + 1;"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 8.98μs -> 5.98μs (50.1% faster)

def test_find_junit_assertions_assertion_without_semicolon():
    """Test assertion statement without trailing semicolon."""
    source = "assertEquals(5, getValue())"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.9μs -> 15.4μs (42.4% faster)

def test_find_junit_assertions_multiline_assertion():
    """Test assertion spanning multiple lines."""
    source = """assertEquals(
        5,
        getValue()
    );"""
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 32.7μs -> 25.4μs (29.0% faster)

def test_find_junit_assertions_nested_parentheses():
    """Test assertion with nested parentheses in arguments."""
    source = "assertEquals(getValue(nested(param)), actual);"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 29.5μs -> 21.5μs (37.0% faster)

def test_find_junit_assertions_string_with_quotes():
    """Test assertion with string literals containing quotes."""
    source = 'assertEquals("expected", getValue("param\\"with\\"quotes"));\n'
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 34.3μs -> 25.7μs (33.5% faster)

def test_find_junit_assertions_string_with_parenthesis():
    """Test assertion with string literal containing parenthesis."""
    source = 'assertEquals("test(value)", getValue());\n'
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 25.7μs -> 18.9μs (36.1% faster)

def test_find_junit_assertions_char_literal():
    """Test assertion with character literal."""
    source = "assertEquals('c', getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 22.2μs -> 16.0μs (39.1% faster)

def test_find_junit_assertions_unmatched_parentheses():
    """Test assertion with unmatched parentheses (malformed)."""
    source = "assertEquals(5, getValue();"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 11.8μs -> 7.74μs (52.3% faster)

def test_find_junit_assertions_commented_out():
    """Test that commented-out assertions are still found by regex."""
    source = "// assertEquals(5, getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.9μs -> 15.5μs (41.4% faster)

def test_find_junit_assertions_in_block_comment():
    """Test that assertions in block comments are still matched by regex."""
    source = "/* assertEquals(5, getValue()); */"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 22.1μs -> 15.6μs (41.6% faster)

def test_find_junit_assertions_assertArrayEquals():
    """Test finding assertArrayEquals assertion."""
    source = "assertArrayEquals(expected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 20.8μs -> 14.4μs (45.1% faster)

def test_find_junit_assertions_assertNotEquals():
    """Test finding assertNotEquals assertion."""
    source = "assertNotEquals(notExpected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 22.0μs -> 15.3μs (43.6% faster)

def test_find_junit_assertions_assertSame():
    """Test finding assertSame assertion."""
    source = "assertSame(expected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 20.7μs -> 13.9μs (48.9% faster)

def test_find_junit_assertions_assertNotSame():
    """Test finding assertNotSame assertion."""
    source = "assertNotSame(notExpected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.4μs -> 14.9μs (43.9% faster)

def test_find_junit_assertions_assertIterableEquals():
    """Test finding assertIterableEquals assertion."""
    source = "assertIterableEquals(expected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 20.2μs -> 13.9μs (45.1% faster)

def test_find_junit_assertions_assertLinesMatch():
    """Test finding assertLinesMatch assertion."""
    source = "assertLinesMatch(expected, actual);"
    transformer = JavaAssertTransformer("actual")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 20.3μs -> 13.9μs (46.1% faster)

def test_find_junit_assertions_assertTimeout():
    """Test finding assertTimeout assertion."""
    source = "assertTimeout(Duration.ofSeconds(2), () -> getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 31.8μs -> 24.4μs (30.3% faster)

def test_find_junit_assertions_assertTimeoutPreemptively():
    """Test finding assertTimeoutPreemptively assertion."""
    source = "assertTimeoutPreemptively(Duration.ofSeconds(2), () -> getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 31.8μs -> 24.2μs (31.3% faster)

def test_find_junit_assertions_assertAll():
    """Test finding assertAll assertion."""
    source = "assertAll(() -> assertEquals(5, getValue()));"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 38.8μs -> 28.1μs (38.2% faster)

def test_find_junit_assertions_with_message():
    """Test assertion with message parameter."""
    source = 'assertEquals(5, getValue(), "Should be 5");'
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 26.5μs -> 19.3μs (37.2% faster)

def test_find_junit_assertions_adjacent_assertions():
    """Test multiple assertions on adjacent lines."""
    source = """assertEquals(5, getValue());
assertTrue(isValid());"""
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 30.3μs -> 21.3μs (42.0% faster)

def test_find_junit_assertions_assertion_in_loop():
    """Test assertion inside a loop."""
    source = """for (int i = 0; i < 10; i++) {
    assertEquals(i, getValue());
}"""
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 24.6μs -> 18.1μs (35.6% faster)

def test_find_junit_assertions_assertion_in_if():
    """Test assertion inside if statement."""
    source = """if (isValid()) {
    assertEquals(5, getValue());
}"""
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 23.2μs -> 16.6μs (39.9% faster)

def test_find_junit_assertions_multiple_target_calls():
    """Test assertion with multiple target function calls."""
    source = "assertEquals(getValue(), getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 26.2μs -> 19.1μs (36.6% faster)

def test_find_junit_assertions_static_method_call():
    """Test target call as static method."""
    source = "assertEquals(5, MyClass.getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 23.6μs -> 16.8μs (40.5% faster)

def test_find_junit_assertions_chained_method_calls():
    """Test target call with chained method calls."""
    source = "assertEquals(5, obj.getBuilder().getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 28.2μs -> 20.9μs (35.2% faster)

def test_find_junit_assertions_lambda_in_assertThrows():
    """Test lambda body extraction in assertThrows."""
    source = "assertThrows(Exception.class, () -> getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 34.1μs -> 25.5μs (33.7% faster)

def test_find_junit_assertions_block_lambda_in_assertThrows():
    """Test block lambda body extraction in assertThrows."""
    source = "assertThrows(Exception.class, () -> { getValue(); });"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 37.4μs -> 28.1μs (33.1% faster)

def test_find_junit_assertions_position_tracking():
    """Test that start and end positions are correctly tracked."""
    source = "assertEquals(5, getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.7μs -> 14.9μs (45.8% faster)

def test_find_junit_assertions_position_with_leading_spaces():
    """Test position tracking with leading whitespace."""
    source = "        assertEquals(5, getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.5μs -> 15.0μs (43.1% faster)
    extracted = source[results[0].start_pos:results[0].end_pos]

def test_find_junit_assertions_no_target_function_match():
    """Test assertion with no matching target function call."""
    source = "assertEquals(5, someOtherFunction());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 22.7μs -> 15.9μs (42.6% faster)

def test_find_junit_assertions_case_sensitive_method():
    """Test that method names are case-sensitive."""
    source = "assertEquals(5, GetValue());"  # GetValue != getValue
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 18.8μs -> 12.9μs (46.0% faster)

def test_find_junit_assertions_many_assertions():
    """Test finding many assertions in a single source."""
    lines = [f"assertEquals({i}, getValue());" for i in range(100)]
    source = "\n".join(lines)
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 1.01ms -> 740μs (37.0% faster)
    for i, result in enumerate(results):
        pass

def test_find_junit_assertions_large_method_body():
    """Test finding assertions in a large method body."""
    # Create a method body with mixed code and assertions
    lines = []
    for i in range(500):
        if i % 5 == 0:
            lines.append(f"    assertEquals({i}, getValue());")
        else:
            lines.append(f"    int x{i} = {i};")
    
    source = "\n".join(lines)
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 1.98ms -> 1.70ms (16.5% faster)

def test_find_junit_assertions_deeply_nested_parentheses():
    """Test assertion with deeply nested parentheses."""
    # Create deeply nested function calls
    nested = "getValue()"
    for _ in range(50):
        nested = f"wrap({nested})"
    
    source = f"assertEquals(5, {nested});"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 122μs -> 96.7μs (27.2% faster)

def test_find_junit_assertions_long_assertion_chain():
    """Test multiple assertions with complex nesting."""
    source = """
    assertEquals(
        functionCall(
            anotherCall(param1, param2),
            getValue()
        ),
        expectedValue
    );
    assertTrue(getValue() != null);
    assertFalse(getValue().isEmpty());
    """
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 101μs -> 83.6μs (21.4% faster)

def test_find_junit_assertions_mixed_frameworks():
    """Test source with mixed assertion styles."""
    source = """
    assertEquals(5, getValue());
    Assert.assertEquals(5, getValue());
    Assertions.assertEquals(5, getValue());
    assertTrue(getValue() > 0);
    AssertJ might appear but shouldn't be matched here;
    """
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 60.4μs -> 45.5μs (32.6% faster)

def test_find_junit_assertions_all_assertion_types():
    """Test finding all types of JUnit assertions."""
    assertions_to_test = [
        "assertEquals(5, getValue());",
        "assertNotEquals(5, getValue());",
        "assertSame(obj, getValue());",
        "assertNotSame(obj, getValue());",
        "assertArrayEquals(arr, getValue());",
        "assertIterableEquals(list, getValue());",
        "assertLinesMatch(lines, getValue());",
        "assertTrue(getValue() > 0);",
        "assertFalse(getValue() < 0);",
        "assertNull(getValue());",
        "assertNotNull(getValue());",
        "assertThrows(Exception.class, () -> getValue());",
        "assertDoesNotThrow(() -> getValue());",
        "assertTimeout(Duration.ofSeconds(1), () -> getValue());",
        "assertTimeoutPreemptively(Duration.ofSeconds(1), () -> getValue());",
        "assertAll(() -> assertEquals(1, getValue()));",
    ]
    
    source = "\n".join(assertions_to_test)
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 222μs -> 171μs (30.1% faster)

def test_find_junit_assertions_performance_large_source():
    """Test performance with large source file."""
    # Create a large source with many different types of statements
    lines = []
    for i in range(500):
        if i % 10 == 0:
            lines.append(f"    assertEquals({i}, getValue());")
        elif i % 10 == 1:
            lines.append(f"    assertTrue(isValid());")
        else:
            lines.append(f"    int var{i} = {i};")
            lines.append(f"    String str{i} = \"value{i}\";")
    
    source = "\n".join(lines)
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 3.14ms -> 2.89ms (8.68% faster)
    
    # Should find assertEquals assertions (50 of them)
    assertEquals_results = [r for r in results if r.assertion_method == "assertEquals"]

def test_find_junit_assertions_unicode_characters():
    """Test assertions with unicode characters in strings."""
    source = 'assertEquals("test\\u00E9", getValue()); // Contains accented character'
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 28.7μs -> 21.9μs (31.0% faster)

def test_find_junit_assertions_escaped_quotes():
    """Test assertions with various escaped quote patterns."""
    source = r'''
    assertEquals("test\"quote", getValue());
    assertEquals('test\'quote', getValue());
    assertEquals("test\\backslash", getValue());
    '''
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 57.4μs -> 44.9μs (28.0% faster)

def test_find_junit_assertions_position_accuracy():
    """Test that positions are accurate across many assertions."""
    source = """assertEquals(1, getValue());
assertEquals(2, getValue());
assertEquals(3, getValue());"""
    
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 43.5μs -> 31.6μs (37.6% faster)
    
    # Verify each position range is non-overlapping and sequential
    prev_end = 0
    for result in results:
        extracted = source[result.start_pos:result.end_pos]
        prev_end = result.end_pos

def test_find_junit_assertions_whitespace_variations():
    """Test assertions with various whitespace patterns."""
    source = """
    assertEquals(5,getValue());
    assertEquals(5 , getValue());
    assertEquals( 5 , getValue() );
    assertEquals(
        5,
        getValue()
    );
    """
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 65.7μs -> 50.3μs (30.7% faster)

def test_find_junit_assertions_no_false_positives_similar_names():
    """Test that similar method names don't create false matches."""
    source = """
    assertEqualsValue(5, getValue());
    someAssertEquals(5, getValue());
    assertEquals(5, getValue());
    """
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 29.9μs -> 23.3μs (28.3% faster)
    
    # Should only find the actual assertEquals
    assertEquals_count = sum(1 for r in results if r.assertion_method == "assertEquals")

def test_find_junit_assertions_receiver_variations():
    """Test extracting target calls with various receiver patterns."""
    source = """
    assertEquals(5, getValue());
    assertEquals(5, obj.getValue());
    assertEquals(5, MyClass.getValue());
    assertEquals(5, obj.inner.getValue());
    """
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 58.8μs -> 43.9μs (33.9% faster)

def test_find_junit_assertions_result_object_attributes():
    """Test that AssertionMatch objects have all required attributes."""
    source = "assertEquals(5, getValue());"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 21.4μs -> 14.7μs (46.2% faster)
    result = results[0]

def test_find_junit_assertions_target_call_object_attributes():
    """Test that TargetCall objects have all required attributes."""
    source = "assertEquals(5, obj.getValue(arg1, arg2));"
    transformer = JavaAssertTransformer("getValue")
    codeflash_output = transformer._find_junit_assertions(source); results = codeflash_output # 27.4μs -> 19.4μs (41.1% faster)
    
    target_call = results[0].target_calls[0]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1295-2026-02-03T21.43.34 and push.

[KRRT7]

Closing stale bot PR.