From 530a6215b94cce8614e49c979dbc211de9533bfc Mon Sep 17 00:00:00 2001
From: Shivani Sharma <s.happyrose@gmail.com>
Date: Fri, 28 Jun 2024 02:57:12 +1000
Subject: [PATCH] Add support for multiple statements in Verifications to be
 migrated by replacing the Verification block

---
 .../java/testing/jmockit/JMockitBlockRewriter.java       | 9 ++++++++-
 .../jmockit/JMockitVerificationsToMockitoTest.java       | 4 +++-
 2 files changed, 11 insertions(+), 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
index 47b1bc3c5..cbe237b7b 100644
--- 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
index 0819adde5..e7492a700 100644
--- 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());