Add support for multiple statements in Verifications to be migrated b…

…y replacing the Verification block
openrewrite · Jun 27, 2024 · 530a621 · 530a621
1 parent 19bc094
commit 530a621
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockRewriter.java b/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockRewriter.java
@@ -26,6 +26,8 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import static org.openrewrite.java.testing.jmockit.JMockitBlockType.Verifications;
+
 class JMockitBlockRewriter {
 
     private static final String WHEN_TEMPLATE_PREFIX = "when(#{any()}).";
@@ -200,14 +202,19 @@ private void writeMethodVerification(J.MethodInvocation invocation, @Nullable Ex
 
         String verifyTemplate = getVerifyTemplate(invocation.getArguments(), fqn, verificationMode, templateParams);
         JavaCoordinates verifyCoordinates;
-        if (this.blockType == JMockitBlockType.Verifications) {
+        if (this.blockType == Verifications) {
             // for Verifications, replace the Verifications block
             verifyCoordinates = nextStatementCoordinates;
+            numStatementsAdded++;
         } else {
             // for Expectations put the verify at the end of the method
             verifyCoordinates = methodBody.getCoordinates().lastStatement();
         }
+
         methodBody = rewriteTemplate(verifyTemplate, templateParams, verifyCoordinates);
+        if (this.blockType == Verifications) {
+            nextStatementCoordinates = this.methodBody.getStatements().get(bodyStatementIndex + numStatementsAdded - 1).getCoordinates().after();
+        }
     }
 
     private J.Block rewriteTemplate(String verifyTemplate, List<Object> templateParams, JavaCoordinates rewriteCoords) {

diff --git a/src/test/java/org/openrewrite/java/testing/jmockit/JMockitVerificationsToMockitoTest.java b/src/test/java/org/openrewrite/java/testing/jmockit/JMockitVerificationsToMockitoTest.java
@@ -726,7 +726,7 @@ void test() {
     }
 
     @Test
-    void whenMultipleVerifications() {
+    void whenMultipleVerificationsAndMultipleStatements() {
         //language=java
         rewriteRun(
           java(
@@ -745,6 +745,7 @@ void test() {
                       myObject.wait();
                       new Verifications() {{
                           myObject.wait();
+                          myObject.wait(anyLong, anyInt);
                       }};
                       myObject.wait(1L);
                       myObject.wait(2L);
@@ -770,6 +771,7 @@ class MyTest {
                   void test() {
                       myObject.wait();
                       verify(myObject).wait();
+                      verify(myObject).wait(anyLong(), anyInt());
                       myObject.wait(1L);
                       myObject.wait(2L);
                       verify(myObject, times(2)).wait(anyLong());