From e88fe26110afd6641c464cb0e87bd898b41bf457 Mon Sep 17 00:00:00 2001
From: Rafael Winterhalter <rafael.wth@gmail.com>
Date: Thu, 11 Feb 2021 20:12:41 +0100
Subject: [PATCH] Add API for clearing mocks. (#2194)

By clearing mocks, caches are emptied and instrumentations are reversed.
---
 src/main/java/org/mockito/Mockito.java        | 11 +++++++++
 .../org/mockito/internal/MockitoCore.java     |  5 ++++
 .../bytebuddy/ByteBuddyMockMaker.java         |  5 ++++
 .../creation/bytebuddy/BytecodeGenerator.java |  2 ++
 .../bytebuddy/InlineByteBuddyMockMaker.java   |  6 +++++
 .../bytebuddy/InlineBytecodeGenerator.java    | 23 ++++++++++++++++++
 .../bytebuddy/SubclassByteBuddyMockMaker.java |  5 ++++
 .../TypeCachingBytecodeGenerator.java         | 18 ++++++++++++++
 .../org/mockito/internal/util/MockUtil.java   |  4 ++++
 .../java/org/mockito/plugins/MockMaker.java   |  6 +++++
 .../java/org/mockito/MockitoClearTest.java    | 24 +++++++++++++++++++
 .../creation/AndroidByteBuddyMockMaker.java   |  5 ++++
 12 files changed, 114 insertions(+)
 create mode 100644 src/test/java/org/mockito/MockitoClearTest.java

diff --git a/src/main/java/org/mockito/Mockito.java b/src/main/java/org/mockito/Mockito.java
index 9d20b43bae..336251e0ae 100644
--- 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
index 6901baef85..cebc40a2e6 100644
--- 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
index 5263ea1262..8153f526ea 100644
--- 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
index b27b9727a1..392392d2d9 100644
--- 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
index 9c0aa9362f..32e5b99f23 100644
--- 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
index 7db8afe97f..78c9d440a1 100644
--- 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
index e112d9e281..8bf72feb89 100644
--- 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
index b3e70c7da7..f7c6703cb2 100644
--- 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
index 25a29d2c0d..2927d6aabf 100644
--- 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
index fcebbb3ae7..7e06e32e5f 100644
--- 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
new file mode 100644
index 0000000000..26b7f76d91
--- /dev/null
+++ 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/subprojects/android/src/main/java/org/mockito/android/internal/creation/AndroidByteBuddyMockMaker.java b/subprojects/android/src/main/java/org/mockito/android/internal/creation/AndroidByteBuddyMockMaker.java
index db77fb7b03..98ffbcb34c 100644
--- a/subprojects/android/src/main/java/org/mockito/android/internal/creation/AndroidByteBuddyMockMaker.java
+++ b/subprojects/android/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();
+    }
 }