Handle multiple blocks in JMockit Expectations

src/main/java/org/openrewrite/java/testing/jmockit/JMockitUtils.java

@@ -36,14 +36,11 @@ static Optional<JMockitBlockType> getJMockitBlock(Statement s) {
         J.Identifier clazz = (J.Identifier) nc.getClazz();
 
         // JMockit block should be composed of a block within another block
-        if (nc.getBody() == null || nc.getBody().getStatements().size() != 1) {
+        if (nc.getBody() == null || (nc.getBody().getStatements().size() != 1 &&
+                !TypeUtils.isAssignableTo("mockit.Expectations", clazz.getType()))) {
             return empty();
         }
 
-        JMockitBlockType blockType = JMockitBlockType.valueOf(clazz.getSimpleName());
-        if (blockType != null && TypeUtils.isOfClassType(clazz.getType(), blockType.getFqn())) {
-            return Optional.of(blockType);
-        }
-        return empty();
+        return Optional.of(JMockitBlockType.valueOf(clazz.getSimpleName()));
     }
 }

src/main/java/org/openrewrite/java/testing/jmockit/SetupStatementsRewriter.java

@@ -52,18 +52,29 @@ J.Block rewriteMethodBody() {
             assert nc.getBody() != null;
             J.Block expectationsBlock = (J.Block) nc.getBody().getStatements().get(0);
 
+            // Account for Expectations which may contain multiple blocks
+            List<Statement> statementList = new ArrayList<>();
+            if (TypeUtils.isAssignableTo("mockit.Expectations", nc.getType())) {
+                statementList.addAll(nc.getBody().getStatements());
+            } else {
+                statementList.add(expectationsBlock);
+            }
+
             // statement needs to be moved directly before expectations class instantiation
             JavaCoordinates coordinates = nc.getCoordinates().before();
             List<Statement> newExpectationsBlockStatements = new ArrayList<>();
-            for (Statement expectationStatement : expectationsBlock.getStatements()) {
-                if (!isSetupStatement(expectationStatement, spies)) {
-                    newExpectationsBlockStatements.add(expectationStatement);
-                    continue;
+            for (Statement st : statementList) {
+                for (Statement expectationStatement : ((J.Block) st).getStatements()) {
+                    if (!isSetupStatement(expectationStatement, spies)) {
+                        newExpectationsBlockStatements.add(expectationStatement);
+                        continue;
+                    }
+                    rewriteBodyStatement(expectationStatement, coordinates);
+                    // subsequent setup statements are moved in order
+                    coordinates = expectationStatement.getCoordinates().after();
                 }
-                rewriteBodyStatement(expectationStatement, coordinates);
-                // subsequent setup statements are moved in order
-                coordinates = expectationStatement.getCoordinates().after();
             }
+
             // the new expectations block has the setup statements removed
             J.Block newExpectationsBlock = expectationsBlock.withStatements(newExpectationsBlockStatements);
             nc = nc.withBody(nc.getBody().withStatements(Collections.singletonList(newExpectationsBlock)));
@@ -128,7 +139,7 @@ private static boolean isNotMockIdentifier(J.Identifier identifier, Set<String>
             return false;
         }
         if (identifier.getType() instanceof JavaType.Method
-                && TypeUtils.isAssignableTo("mockit.Invocations",
+            && TypeUtils.isAssignableTo("mockit.Invocations",
                 ((JavaType.Method) identifier.getType()).getDeclaringType())) {
             return false;
         }
@@ -138,8 +149,8 @@ private static boolean isNotMockIdentifier(J.Identifier identifier, Set<String>
         }
         for (JavaType.FullyQualified annotationType : fieldType.getAnnotations()) {
             if (TypeUtils.isAssignableTo("mockit.Mocked", annotationType)
-                    || TypeUtils.isAssignableTo("mockit.Injectable", annotationType)
-                    || TypeUtils.isAssignableTo("mockit.Tested", annotationType)) {
+                || TypeUtils.isAssignableTo("mockit.Injectable", annotationType)
+                || TypeUtils.isAssignableTo("mockit.Tested", annotationType)) {
                 return false;
             }
         }

src/test/java/org/openrewrite/java/testing/jmockit/JMockitExpectationsToMockitoTest.java

@@ -1419,7 +1419,80 @@ void test() {
     }
 
     @Test
-    void whenMultipleExpectationsNoResults() {
+    void whenMultipleBlockInSingleExpectation() {
+        //language=java
+        rewriteRun(
+          java(
+            """
+              class MyObject {
+                  public String getX() {
+                      return "X";
+                  }
+                  public String getY() {
+                      return "Y";
+                  }
+              }
+              """
+          ),
+          java(
+            """
+              import mockit.Expectations;
+              import mockit.Mocked;
+              import mockit.integration.junit5.JMockitExtension;
+              import org.junit.jupiter.api.extension.ExtendWith;
+
+              import static org.junit.jupiter.api.Assertions.assertEquals;
+              import static org.junit.jupiter.api.Assertions.assertNull;
+
+              @ExtendWith(JMockitExtension.class)
+              class MyTest {
+                  @Mocked
+                  MyObject myObject;
+
+                  void test() {
+                      new Expectations() {
+                        { 
+                          myObject.getX();
+                          result = "x1";
+                        }
+                        {
+                          myObject.getY();
+                          result = "y1";
+                        }
+                      };
+                      assertEquals("x1", myObject.getX());
+                      assertEquals("y1", myObject.getY());
+                  }
+              }
+              """,
+            """
+              import org.junit.jupiter.api.extension.ExtendWith;
+              import org.mockito.Mock;
+              import org.mockito.junit.jupiter.MockitoExtension;
+
+              import static org.junit.jupiter.api.Assertions.assertEquals;
+              import static org.junit.jupiter.api.Assertions.assertNull;
+              import static org.mockito.Mockito.when;
+
+              @ExtendWith(MockitoExtension.class)
+              class MyTest {
+                  @Mock
+                  MyObject myObject;
+
+                  void test() {
+                      when(myObject.getX()).thenReturn("x1");
+                      when(myObject.getY()).thenReturn("y1");
+                      assertEquals("x1", myObject.getX());
+                      assertEquals("y1", myObject.getY());
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void whenMultipleExpectationsNoResults() { 
         //language=java
         rewriteRun(
           java(