fix Mockito memory leak

Author: CoolSpy3

src/main/java/frc/robot/lib/DummySparkMaxAnswer.java

@@ -9,7 +9,6 @@
 import com.revrobotics.CANSparkMaxLowLevel.MotorType;
 import com.revrobotics.SparkMaxPIDController.AccelStrategy;
 
-import org.mockito.Mockito;
 import org.mockito.invocation.InvocationOnMock;
 
 public class DummySparkMaxAnswer extends REVLibErrorAnswer {
@@ -18,12 +17,12 @@ public class DummySparkMaxAnswer extends REVLibErrorAnswer {
 
     public static final DummySparkMaxAnswer ANSWER = new DummySparkMaxAnswer();
 
-    public static final CANSparkMax DUMMY_SPARK_MAX = Mockito.mock(CANSparkMax.class, ANSWER);
+    public static final CANSparkMax DUMMY_SPARK_MAX = Mocks.mock(CANSparkMax.class, ANSWER);
 
-    public static final RelativeEncoder DUMMY_ENCODER = Mockito.mock(RelativeEncoder.class, REVLibErrorAnswer.ANSWER);
-    public static final SparkMaxAnalogSensor DUMMY_ANALOG_SENSOR = Mockito.mock(SparkMaxAnalogSensor.class, REVLibErrorAnswer.ANSWER);
-    public static final SparkMaxLimitSwitch DUMMY_LIMIT_SWITCH = Mockito.mock(SparkMaxLimitSwitch.class, REVLibErrorAnswer.ANSWER);
-    public static final SparkMaxPIDController DUMMY_PID_CONTROLLER = Mockito.mock(SparkMaxPIDController.class, ANSWER);
+    public static final RelativeEncoder DUMMY_ENCODER = Mocks.mock(RelativeEncoder.class, REVLibErrorAnswer.ANSWER);
+    public static final SparkMaxAnalogSensor DUMMY_ANALOG_SENSOR = Mocks.mock(SparkMaxAnalogSensor.class, REVLibErrorAnswer.ANSWER);
+    public static final SparkMaxLimitSwitch DUMMY_LIMIT_SWITCH = Mocks.mock(SparkMaxLimitSwitch.class, REVLibErrorAnswer.ANSWER);
+    public static final SparkMaxPIDController DUMMY_PID_CONTROLLER = Mocks.mock(SparkMaxPIDController.class, ANSWER);
 
     @Override
     public Object answer(InvocationOnMock invocation) throws Throwable {

src/main/java/frc/robot/lib/Mocks.java

@@ -1,20 +1,44 @@
 package frc.robot.lib;
 
+import java.lang.ref.WeakReference;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
-import org.mockito.MockSettings;
 
+import org.mockito.MockSettings;
 import org.mockito.Mockito;
 import org.mockito.internal.stubbing.defaultanswers.ReturnsSmartNulls;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
 public final class Mocks {
+
+    private static final List<WeakReference<Object>> MOCKS = Collections.synchronizedList(new ArrayList<>());
+    private static final Predicate<WeakReference<?>> IS_REFERENCE_CLEARED = reference -> reference.get() == null;
+    private static final Consumer<WeakReference<Object>> CLEAR_INVOCATIONS_ON_REFERENCED_MOCK = reference -> Mockito.clearInvocations(reference.get());
+    private static final Predicate<WeakReference<Object>> CLEAR_INVOCATIONS_ON_REFERENCED_MOCK_IF_REFERNCE_NOT_CLEARED = reference -> {
+        if(IS_REFERENCE_CLEARED.test(reference)) return true;
+        CLEAR_INVOCATIONS_ON_REFERENCED_MOCK.accept(reference);
+        return false;
+    };
+
+    static {
+        // Use a single predicate so that clearing references and invocations is an atomic operation
+        // Otherwise, we could (rarely) run into:
+        // 1) Mock is added
+        // 2) Garbage collected references are removed
+        // 3) Mock is garbage collected
+        // 4) Mock invocations are cleared -> throws NullPointerException
+        Lib199Subsystem.registerPeriodic(() -> MOCKS.removeIf(CLEAR_INVOCATIONS_ON_REFERENCED_MOCK_IF_REFERNCE_NOT_CLEARED));
+    }
 
     /**
      * Attempts to create an instance of a class in which some or all of the classes methods are replaced with a mocked implementation
@@ -77,7 +101,7 @@ public static <T, U> T createMock(Class<T> classToMock, U implClass, Answer<Obje
             settings = Mockito.withSettings().extraInterfaces(interfaces);
         }
         settings = settings.defaultAnswer(new MockAnswer<>(methods, implClass, defaultAnswer));
-        T mock = Mockito.mock(classToMock, settings);
+        T mock = mock(classToMock, settings);
         return mock;
     }
 
@@ -88,6 +112,55 @@ public static Method[] listMethods(Class<?> base, Class<?>... interfaces) {
         return out.toArray(Method[]::new);
     }
 
+    /**
+     * A wrapper for the underlying Mockito method which automatically calls {@link Mockito#clearInvocations(Object...)} to prevent memory leaks
+     * 
+     * @see Mockito#mock(Class)
+     */
+    public static <T> T mock(Class<T> classToMock) {
+        return reportMock(Mockito.mock(classToMock));
+    }
+
+    /**
+     * A wrapper for the underlying Mockito method which automatically calls {@link Mockito#clearInvocations(Object...)} to prevent memory leaks
+     * 
+     * @see Mockito#mock(Class, String)
+     */
+    public static <T> T mock(Class<T> classToMock, String name) {
+        return reportMock(Mockito.mock(classToMock, name));
+    }
+
+    /**
+     * A wrapper for the underlying Mockito method which automatically calls {@link Mockito#clearInvocations(Object...)} to prevent memory leaks
+     * 
+     * @see Mockito#mock(Class, Answer)
+     */
+    public static <T> T mock(Class<T> classToMock, Answer<?> defaultAnswer) {
+        return reportMock(Mockito.mock(classToMock, defaultAnswer));
+    }
+
+    /**
+     * A wrapper for the underlying Mockito method which automatically calls {@link Mockito#clearInvocations(Object...)} to prevent memory leaks
+     * 
+     * @see Mockito#mock(Class, MockSettings)
+     */
+    public static <T> T mock(Class<T> classToMock, MockSettings mockSettings) {
+        return reportMock(Mockito.mock(classToMock, mockSettings));
+    }
+
+    /**
+     * Registers a Mockito mock and periodically calls {@link Mockito#clearInvocations(Object...)} on it to prevent memory leaks
+     * 
+     * @param <T> The type of the mock
+     * @param t The mock
+     * @return The mock
+     */
+    public static <T> T reportMock(T t) {
+        // Wrap in a WeakReference to prevent memory leaks on objects with no more references
+        if(Mockito.mockingDetails(t).isMock()) MOCKS.add(new WeakReference<Object>(t));
+        return t;
+    }
+
     private Mocks() {}
 
     private static final class MockAnswer<U> implements Answer<Object> {

src/test/java/frc/robot/lib/DummySparkMaxAnswerTest.java

@@ -12,12 +12,11 @@
 import com.revrobotics.SparkMaxPIDController.AccelStrategy;
 
 import org.junit.Test;
-import org.mockito.Mockito;
 
 public class DummySparkMaxAnswerTest {
 
     public CANSparkMax createMockedSparkMax() {
-        return Mockito.mock(CANSparkMax.class, new DummySparkMaxAnswer());
+        return Mocks.mock(CANSparkMax.class, new DummySparkMaxAnswer());
     }
 
     public static void assertTestResponses(CANSparkMax spark) {

src/test/java/frc/robot/lib/ErrorCodeAnswerTest.java

@@ -6,13 +6,12 @@
 import com.ctre.phoenix.motorcontrol.can.WPI_TalonSRX;
 
 import org.junit.Test;
-import org.mockito.Mockito;
 
 public class ErrorCodeAnswerTest {
 
     @Test
     public void testResponses() throws Exception {
-        WPI_TalonSRX talon = Mockito.mock(WPI_TalonSRX.class, new ErrorCodeAnswer());
+        WPI_TalonSRX talon = Mocks.mock(WPI_TalonSRX.class, new ErrorCodeAnswer());
 
         // Check that primative types return "null"
         assertEquals(0, talon.get(), 0.01);

src/test/java/frc/robot/lib/REVLibErrorAnswerTest.java

@@ -6,13 +6,12 @@
 import com.revrobotics.RelativeEncoder;
 
 import org.junit.Test;
-import org.mockito.Mockito;
 
 public class REVLibErrorAnswerTest {
 
     @Test
     public void testResponses() throws Exception {
-        RelativeEncoder enc = Mockito.mock(RelativeEncoder.class, new REVLibErrorAnswer());
+        RelativeEncoder enc = Mocks.mock(RelativeEncoder.class, new REVLibErrorAnswer());
 
         // Check that primative types return "null"
         assertEquals(0, enc.getPosition(), 0.01);