Add API for clearing mocks. (#2194)

By clearing mocks, caches are emptied and instrumentations are reversed.
mockito · Feb 11, 2021 · e88fe26 · e88fe26
1 parent 3bffcd0
commit e88fe26
Show file tree

Hide file tree

Showing 12 changed files with 114 additions and 0 deletions.
diff --git a/src/main/java/org/mockito/Mockito.java b/src/main/java/org/mockito/Mockito.java
@@ -2461,6 +2461,17 @@ public static <T> void reset(T... mocks) {
         MOCKITO_CORE.reset(mocks);
     }
 
+    /**
+     * Clears all mocks, type caches and instrumentations.
+     * <p>
+     * By clearing Mockito's state, previously created mocks might begin to malfunction. This option can be used if
+     * Mockito's caches take up too much space or if the inline mock maker's instrumentation is causing performance
+     * issues in code where mocks are no longer used. Normally, you would not need to use this option.
+     */
+    public static void clearAllCaches() {
+        MOCKITO_CORE.clearAllCaches();
+    }
+
     /**
      * Use this method in order to only clear invocations, when stubbing is non-trivial. Use-cases can be:
      * <ul>

diff --git a/src/main/java/org/mockito/internal/MockitoCore.java b/src/main/java/org/mockito/internal/MockitoCore.java
@@ -15,6 +15,7 @@
 import org.mockito.internal.stubbing.OngoingStubbingImpl;
 import org.mockito.internal.stubbing.StubberImpl;
 import org.mockito.internal.util.DefaultMockingDetails;
+import org.mockito.internal.util.MockUtil;
 import org.mockito.internal.verification.MockAwareVerificationMode;
 import org.mockito.internal.verification.VerificationDataImpl;
 import org.mockito.internal.verification.VerificationModeFactory;
@@ -279,4 +280,8 @@ public MockingDetails mockingDetails(Object toInspect) {
     public LenientStubber lenient() {
         return new DefaultLenientStubber();
     }
+
+    public void clearAllCaches() {
+        MockUtil.clearAllCaches();
+    }
 }
diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/ByteBuddyMockMaker.java b/src/main/java/org/mockito/internal/creation/bytebuddy/ByteBuddyMockMaker.java
@@ -71,4 +71,9 @@ public <T> ConstructionMockControl<T> createConstructionMock(
         return defaultByteBuddyMockMaker.createConstructionMock(
                 type, settingsFactory, handlerFactory, mockInitializer);
     }
+
+    @Override
+    public void clearAllCaches() {
+        defaultByteBuddyMockMaker.clearAllCaches();
+    }
 }
diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/BytecodeGenerator.java b/src/main/java/org/mockito/internal/creation/bytebuddy/BytecodeGenerator.java
@@ -11,4 +11,6 @@ public interface BytecodeGenerator {
     void mockClassConstruction(Class<?> type);
 
     void mockClassStatic(Class<?> type);
+
+    default void clearAllCaches() {}
 }
diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/InlineByteBuddyMockMaker.java b/src/main/java/org/mockito/internal/creation/bytebuddy/InlineByteBuddyMockMaker.java
@@ -482,6 +482,12 @@ public void resetMock(Object mock, MockHandler newHandler, MockCreationSettings
         }
     }
 
+    @Override
+    public void clearAllCaches() {
+        clearAllMocks();
+        bytecodeGenerator.clearAllCaches();
+    }
+
     @Override
     public void clearMock(Object mock) {
         if (mock instanceof Class<?>) {

diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/InlineBytecodeGenerator.java b/src/main/java/org/mockito/internal/creation/bytebuddy/InlineBytecodeGenerator.java
@@ -34,6 +34,7 @@
 
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
+import java.lang.instrument.UnmodifiableClassException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.security.ProtectionDomain;
@@ -396,6 +397,28 @@ public byte[] transform(
         }
     }
 
+    @Override
+    public synchronized void clearAllCaches() {
+        Set<Class<?>> types = new HashSet<>();
+        mocked.forEach(types::add);
+        if (types.isEmpty()) {
+            return;
+        }
+        mocked.clear();
+        flatMocked.clear();
+        try {
+            instrumentation.retransformClasses(types.toArray(new Class<?>[0]));
+        } catch (UnmodifiableClassException e) {
+            throw new MockitoException(
+                    join(
+                            "Failed to reset mocks.",
+                            "",
+                            "This should not influence the working of Mockito.",
+                            "But if the reset intends to remove mocking code to improve performance, it is still impacted."),
+                    e);
+        }
+    }
+
     private static class ParameterWritingVisitorWrapper extends AsmVisitorWrapper.AbstractBase {
 
         private final Class<?> type;

diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/SubclassByteBuddyMockMaker.java b/src/main/java/org/mockito/internal/creation/bytebuddy/SubclassByteBuddyMockMaker.java
@@ -179,4 +179,9 @@ public String nonMockableReason() {
             }
         };
     }
+
+    @Override
+    public void clearAllCaches() {
+        cachingMockBytecodeGenerator.clearAllCaches();
+    }
 }
diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/TypeCachingBytecodeGenerator.java b/src/main/java/org/mockito/internal/creation/bytebuddy/TypeCachingBytecodeGenerator.java
@@ -7,6 +7,8 @@
 import java.lang.ref.ReferenceQueue;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import net.bytebuddy.TypeCache;
 import org.mockito.mock.SerializableMode;
@@ -20,6 +22,8 @@ class TypeCachingBytecodeGenerator extends ReferenceQueue<ClassLoader>
 
     private final TypeCache<MockitoMockKey> typeCache;
 
+    private final ReadWriteLock lock = new ReentrantReadWriteLock();
+
     public TypeCachingBytecodeGenerator(BytecodeGenerator bytecodeGenerator, boolean weak) {
         this.bytecodeGenerator = bytecodeGenerator;
         typeCache =
@@ -30,6 +34,7 @@ public TypeCachingBytecodeGenerator(BytecodeGenerator bytecodeGenerator, boolean
     @SuppressWarnings("unchecked")
     @Override
     public <T> Class<T> mockClass(final MockFeatures<T> params) {
+        lock.readLock().lock();
         try {
             ClassLoader classLoader = params.mockedType.getClassLoader();
             return (Class<T>)
@@ -54,6 +59,8 @@ public Class<?> call() throws Exception {
             } else {
                 throw exception;
             }
+        } finally {
+            lock.readLock().unlock();
         }
     }
 
@@ -67,6 +74,17 @@ public void mockClassConstruction(Class<?> type) {
         bytecodeGenerator.mockClassConstruction(type);
     }
 
+    @Override
+    public void clearAllCaches() {
+        lock.writeLock().lock();
+        try {
+            typeCache.clear();
+            bytecodeGenerator.clearAllCaches();
+        } finally {
+            lock.writeLock().unlock();
+        }
+    }
+
     private static class MockitoMockKey extends TypeCache.SimpleKey {
 
         private final SerializableMode serializableMode;

diff --git a/src/main/java/org/mockito/internal/util/MockUtil.java b/src/main/java/org/mockito/internal/util/MockUtil.java
@@ -152,4 +152,8 @@ public static <T> MockMaker.ConstructionMockControl<T> createConstructionMock(
         return mockMaker.createConstructionMock(
                 type, settingsFactory, handlerFactory, mockInitializer);
     }
+
+    public static void clearAllCaches() {
+        mockMaker.clearAllCaches();
+    }
 }
diff --git a/src/main/java/org/mockito/plugins/MockMaker.java b/src/main/java/org/mockito/plugins/MockMaker.java
@@ -198,6 +198,12 @@ default <T> ConstructionMockControl<T> createConstructionMock(
                         "Note that Mockito's inline mock maker is not supported on Android."));
     }
 
+    /**
+     * Clears all cashes for mocked types and removes all byte code alterations, if possible.
+     */
+    @Incubating
+    default void clearAllCaches() {}
+
     /**
      * Carries the mockability information
      *

diff --git a/src/test/java/org/mockito/MockitoClearTest.java b/src/test/java/org/mockito/MockitoClearTest.java
@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2021 Mockito contributors
+ * This program is made available under the terms of the MIT License.
+ */
+package org.mockito;
+
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class MockitoClearTest {
+
+    @Test
+    public void can_clear_mock() {
+        Base mock = Mockito.mock(Base.class);
+        assertThat(Mockito.mock(Base.class).getClass()).isEqualTo(mock.getClass());
+
+        Mockito.clearAllCaches();
+
+        assertThat(Mockito.mock(Base.class).getClass()).isNotEqualTo(mock.getClass());
+    }
+
+    abstract static class Base {}
+}
diff --git a/...ndroid/src/main/java/org/mockito/android/internal/creation/AndroidByteBuddyMockMaker.java b/...ndroid/src/main/java/org/mockito/android/internal/creation/AndroidByteBuddyMockMaker.java
@@ -53,4 +53,9 @@ public void resetMock(Object mock, MockHandler newHandler, MockCreationSettings
     public TypeMockability isTypeMockable(Class<?> type) {
         return delegate.isTypeMockable(type);
     }
+
+    @Override
+    public void clearAllCaches() {
+        delegate.clearAllCaches();
+    }
 }