From 328a2e3dc6805f34b6daa01818c777a1691898bf Mon Sep 17 00:00:00 2001
From: Tim te Beek <tim@moderne.io>
Date: Tue, 2 Jul 2024 12:00:18 +0200
Subject: [PATCH] Add missing nullable annotations; add Lombok annotations

---
 .../testing/jmockit/JMockitBlockRewriter.java | 65 ++++---------------
 .../testing/jmockit/JMockitBlockType.java     | 16 ++---
 2 files changed, 19 insertions(+), 62 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 2a19884f9..e561d8bcf 100644
--- a/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockRewriter.java
+++ b/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockRewriter.java
@@ -15,6 +15,9 @@
  */
 package org.openrewrite.java.testing.jmockit;
 
+import lombok.AccessLevel;
+import lombok.Data;
+import lombok.Setter;
 import org.openrewrite.Cursor;
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.internal.lang.Nullable;
@@ -191,11 +194,6 @@ private void rewriteVerify(J.MethodInvocation invocation, @Nullable Expression t
         }
         templateParams.add(invocation.getName().getSimpleName());
         String verifyTemplate = getVerifyTemplate(invocation.getArguments(), verificationMode, templateParams);
-        if (verifyTemplate == null) {
-            // invalid template, cannot rewrite
-            rewriteFailed = true;
-            return;
-        }
 
         JavaCoordinates verifyCoordinates;
         if (this.blockType == Verifications) {
@@ -244,7 +242,7 @@ private J.Block rewriteTemplate(String verifyTemplate, List<Object> templatePara
                 );
     }
 
-    private static String getWhenTemplate(List<Expression> results) {
+    private static @Nullable String getWhenTemplate(List<Expression> results) {
         boolean buildingResults = false;
         final StringBuilder templateBuilder = new StringBuilder(WHEN_TEMPLATE_PREFIX);
         for (Expression result : results) {
@@ -320,7 +318,7 @@ private static String getVerifyTemplate(List<Expression> arguments, @Nullable St
         return templateBuilder.toString();
     }
 
-    private static MockInvocationResults buildMockInvocationResults(List<Statement> expectationStatements) {
+    private static @Nullable MockInvocationResults buildMockInvocationResults(List<Statement> expectationStatements) {
         final MockInvocationResults resultWrapper = new MockInvocationResults();
         for (int i = 1; i < expectationStatements.size(); i++) {
             Statement expectationStatement = expectationStatements.get(i);
@@ -356,20 +354,19 @@ private static MockInvocationResults buildMockInvocationResults(List<Statement>
         return resultWrapper;
     }
 
-    private static String getVariableNameFromAssignment(J.Assignment assignment) {
-        String name = null;
+    private static @Nullable String getVariableNameFromAssignment(J.Assignment assignment) {
         if (assignment.getVariable() instanceof J.Identifier) {
-            name = ((J.Identifier) assignment.getVariable()).getSimpleName();
+            return ((J.Identifier) assignment.getVariable()).getSimpleName();
         } else if (assignment.getVariable() instanceof J.FieldAccess) {
             J.FieldAccess fieldAccess = (J.FieldAccess) assignment.getVariable();
             if (fieldAccess.getTarget() instanceof J.Identifier) {
-                name = fieldAccess.getSimpleName();
+                return fieldAccess.getSimpleName();
             }
         }
-        return name;
+        return null;
     }
 
-    private static String getPrimitiveTemplateField(JavaType.Primitive primitiveType) {
+    private static @Nullable String getPrimitiveTemplateField(JavaType.Primitive primitiveType) {
         switch (primitiveType) {
             case Boolean:
                 return "#{any(boolean)}";
@@ -396,54 +393,16 @@ private static String getPrimitiveTemplateField(JavaType.Primitive primitiveType
         }
     }
 
-    private static String getInvocationSelectFullyQualifiedClassName(J.MethodInvocation invocation) {
-        Expression select = invocation.getSelect();
-        if (select == null || select.getType() == null) {
-            return null;
-        }
-        String fqn = null;
-        if (select.getType() instanceof JavaType.FullyQualified) {
-            fqn = ((JavaType.FullyQualified) select.getType()).getFullyQualifiedName();
-        }
-        return fqn;
-    }
-
+    @Data
     private static class MockInvocationResults {
+        @Setter(AccessLevel.NONE)
         private final List<Expression> results = new ArrayList<>();
         private Expression times;
         private Expression minTimes;
         private Expression maxTimes;
 
-        private List<Expression> getResults() {
-            return results;
-        }
-
         private void addResult(Expression result) {
             results.add(result);
         }
-
-        private Expression getTimes() {
-            return times;
-        }
-
-        private void setTimes(Expression times) {
-            this.times = times;
-        }
-
-        private Expression getMinTimes() {
-            return minTimes;
-        }
-
-        private void setMinTimes(Expression minTimes) {
-            this.minTimes = minTimes;
-        }
-
-        private Expression getMaxTimes() {
-            return maxTimes;
-        }
-
-        private void setMaxTimes(Expression maxTimes) {
-            this.maxTimes = maxTimes;
-        }
     }
 }
diff --git a/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockType.java b/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockType.java
index 12cd0b634..e649448a0 100644
--- a/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockType.java
+++ b/src/main/java/org/openrewrite/java/testing/jmockit/JMockitBlockType.java
@@ -15,17 +15,15 @@
  */
 package org.openrewrite.java.testing.jmockit;
 
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+@Getter
 enum JMockitBlockType {
 
-    Expectations("mockit.Expectations"), Verifications("mockit.Verifications"); // Add NonStrictExpectations later
+    Expectations("mockit.Expectations"),
+    Verifications("mockit.Verifications"); // Add NonStrictExpectations later
 
     private final String fqn;
-
-    JMockitBlockType(String fqn) {
-        this.fqn = fqn;
-    }
-
-    String getFqn() {
-        return fqn;
-    }
 }