From 2931c7d8d0688b6bd715970c709bf3486a18cf88 Mon Sep 17 00:00:00 2001
From: Stephan Schroevers <stephan.schroevers@teampicnic.com>
Date: Sat, 2 Sep 2023 20:16:30 +0200
Subject: [PATCH] Introduce additional Mockito Refaster rules

---
 .../refasterrules/MockitoRules.java           | 28 +++++++++++++++++++
 .../refasterrules/MockitoRulesTestInput.java  | 11 ++++++++
 .../refasterrules/MockitoRulesTestOutput.java | 11 ++++++++
 3 files changed, 50 insertions(+)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/MockitoRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/MockitoRules.java
index dbb7b16770..10fa10ebdf 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/MockitoRules.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/MockitoRules.java
@@ -5,10 +5,12 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import com.google.errorprone.refaster.Refaster;
 import com.google.errorprone.refaster.annotation.AfterTemplate;
 import com.google.errorprone.refaster.annotation.BeforeTemplate;
 import com.google.errorprone.refaster.annotation.UseImportPolicy;
 import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.verification.VerificationMode;
 import tech.picnic.errorprone.refaster.annotation.OnlineDocumentation;
 
@@ -50,4 +52,30 @@ T after(T mock) {
       return verify(mock);
     }
   }
+
+  static final class InvocationOnMockGetArguments {
+    @BeforeTemplate
+    Object before(InvocationOnMock invocation, int i) {
+      return invocation.getArguments()[i];
+    }
+
+    @AfterTemplate
+    Object after(InvocationOnMock invocation, int i) {
+      return invocation.getArgument(i);
+    }
+  }
+
+  static final class InvocationOnMockGetArgumentsWithTypeParameter<T> {
+    @BeforeTemplate
+    @SuppressWarnings("unchecked")
+    T before(InvocationOnMock invocation, int i) {
+      return Refaster.anyOf(
+          invocation.getArgument(i, Refaster.<T>clazz()), (T) invocation.getArgument(i));
+    }
+
+    @AfterTemplate
+    T after(InvocationOnMock invocation, int i) {
+      return invocation.<T>getArgument(i);
+    }
+  }
 }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestInput.java
index 85dedcc48c..b06b8a9211 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestInput.java
@@ -5,6 +5,7 @@
 import static org.mockito.Mockito.verify;
 
 import com.google.common.collect.ImmutableSet;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.verification.VerificationMode;
 import tech.picnic.errorprone.refaster.test.RefasterRuleCollectionTestCase;
 
@@ -21,4 +22,14 @@ VerificationMode testNever() {
   Object testVerifyOnce() {
     return verify(mock(Object.class), times(1));
   }
+
+  Object testInvocationOnMockGetArguments() {
+    return ((InvocationOnMock) null).getArguments()[0];
+  }
+
+  ImmutableSet<Number> testInvocationOnMockGetArgumentsWithTypeParameter() {
+    return ImmutableSet.of(
+        ((InvocationOnMock) null).getArgument(0, Integer.class),
+        (Double) ((InvocationOnMock) null).getArgument(1));
+  }
 }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestOutput.java
index f24f864531..bebaad0cae 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/MockitoRulesTestOutput.java
@@ -6,6 +6,7 @@
 import static org.mockito.Mockito.verify;
 
 import com.google.common.collect.ImmutableSet;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.verification.VerificationMode;
 import tech.picnic.errorprone.refaster.test.RefasterRuleCollectionTestCase;
 
@@ -22,4 +23,14 @@ VerificationMode testNever() {
   Object testVerifyOnce() {
     return verify(mock(Object.class));
   }
+
+  Object testInvocationOnMockGetArguments() {
+    return ((InvocationOnMock) null).getArgument(0);
+  }
+
+  ImmutableSet<Number> testInvocationOnMockGetArgumentsWithTypeParameter() {
+    return ImmutableSet.of(
+        ((InvocationOnMock) null).<Integer>getArgument(0),
+        ((InvocationOnMock) null).<Double>getArgument(1));
+  }
 }