Fixes mockito#2626 : Introduce MockSettings.mockMaker

JojOatXGME · Jul 4, 2022 · e322fa5 · e322fa5
1 parent e962176
commit e322fa5
Show file tree

Hide file tree

Showing 13 changed files with 138 additions and 17 deletions.
diff --git a/src/main/java/org/mockito/MockSettings.java b/src/main/java/org/mockito/MockSettings.java
@@ -15,6 +15,7 @@
 import org.mockito.listeners.VerificationStartedListener;
 import org.mockito.mock.MockCreationSettings;
 import org.mockito.mock.SerializableMode;
+import org.mockito.plugins.MockMaker;
 import org.mockito.quality.Strictness;
 import org.mockito.stubbing.Answer;
 
@@ -381,4 +382,6 @@ public interface MockSettings extends Serializable {
      * @since 4.6.0
      */
     MockSettings strictness(Strictness strictness);
+
+    MockSettings mockMaker(Class<? extends MockMaker> mockMaker);
 }
diff --git a/src/main/java/org/mockito/internal/MockitoCore.java b/src/main/java/org/mockito/internal/MockitoCore.java
@@ -160,6 +160,14 @@ public <T> MockedConstruction<T> mockConstruction(
                                         + "At the moment, you cannot provide your own implementations of that class.");
                     }
                     MockSettingsImpl impl = MockSettingsImpl.class.cast(value);
+                    Class<? extends MockMaker> mockMaker = impl.getMockMaker();
+                    if (mockMaker != null) {
+                        throw new IllegalArgumentException(
+                                "Unexpected MockMaker '"
+                                        + mockMaker.getCanonicalName()
+                                        + "'\n"
+                                        + "You cannot override the MockMaker for construction mocks.");
+                    }
                     return impl.build(typeToMock);
                 };
         MockMaker.ConstructionMockControl<T> control =

diff --git a/src/main/java/org/mockito/internal/creation/MockSettingsImpl.java b/src/main/java/org/mockito/internal/creation/MockSettingsImpl.java
@@ -34,6 +34,7 @@
 import org.mockito.mock.MockCreationSettings;
 import org.mockito.mock.MockName;
 import org.mockito.mock.SerializableMode;
+import org.mockito.plugins.MockMaker;
 import org.mockito.quality.Strictness;
 import org.mockito.stubbing.Answer;
 
@@ -254,11 +255,17 @@ public MockSettings strictness(Strictness strictness) {
         return this;
     }
 
+    @Override
+    public MockSettings mockMaker(Class<? extends MockMaker> mockMaker) {
+        this.mockMaker = mockMaker;
+        return this;
+    }
+
     private static <T> CreationSettings<T> validatedSettings(
             Class<T> typeToMock, CreationSettings<T> source) {
         MockCreationValidator validator = new MockCreationValidator();
 
-        validator.validateType(typeToMock);
+        validator.validateType(typeToMock, source.getMockMaker());
         validator.validateExtraInterfaces(typeToMock, source.getExtraInterfaces());
         validator.validateMockedType(typeToMock, source.getSpiedInstance());
 

diff --git a/src/main/java/org/mockito/internal/creation/settings/CreationSettings.java b/src/main/java/org/mockito/internal/creation/settings/CreationSettings.java
@@ -18,6 +18,7 @@
 import org.mockito.mock.MockCreationSettings;
 import org.mockito.mock.MockName;
 import org.mockito.mock.SerializableMode;
+import org.mockito.plugins.MockMaker;
 import org.mockito.quality.Strictness;
 import org.mockito.stubbing.Answer;
 
@@ -46,6 +47,7 @@ public class CreationSettings<T> implements MockCreationSettings<T>, Serializabl
     private Object outerClassInstance;
     private Object[] constructorArgs;
     protected Strictness strictness = null;
+    protected Class<? extends MockMaker> mockMaker;
 
     public CreationSettings() {}
 
@@ -68,6 +70,7 @@ public CreationSettings(CreationSettings copy) {
         this.constructorArgs = copy.getConstructorArgs();
         this.strictness = copy.strictness;
         this.stripAnnotations = copy.stripAnnotations;
+        this.mockMaker = copy.mockMaker;
     }
 
     @Override
@@ -178,4 +181,9 @@ public boolean isLenient() {
     public Strictness getStrictness() {
         return strictness;
     }
+
+    @Override
+    public Class<? extends MockMaker> getMockMaker() {
+        return mockMaker;
+    }
 }
diff --git a/.../java/org/mockito/internal/stubbing/defaultanswers/RetrieveGenericsForDefaultAnswers.java b/.../java/org/mockito/internal/stubbing/defaultanswers/RetrieveGenericsForDefaultAnswers.java
@@ -38,6 +38,7 @@ static Object returnTypeForMockWithCorrectGenerics(
         }
 
         if (type != null) {
+            // TODO: Should we use the mockMaker of the mock?
             if (!MOCKITO_CORE.isTypeMockable(type)) {
                 return null;
             }

diff --git a/src/main/java/org/mockito/internal/stubbing/defaultanswers/ReturnsDeepStubs.java b/src/main/java/org/mockito/internal/stubbing/defaultanswers/ReturnsDeepStubs.java
@@ -52,6 +52,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
                         .resolveGenericReturnType(invocation.getMethod());
 
         Class<?> rawType = returnTypeGenericMetadata.rawType();
+        // TODO: Should we use the mockMaker from the mock?
         if (!mockitoCore().isTypeMockable(rawType)) {
             if (invocation.getMethod().getReturnType().equals(rawType)) {
                 return delegate().answer(invocation);

diff --git a/src/main/java/org/mockito/internal/util/MockCreationValidator.java b/src/main/java/org/mockito/internal/util/MockCreationValidator.java
@@ -13,13 +13,14 @@
 import java.util.Collection;
 
 import org.mockito.mock.SerializableMode;
+import org.mockito.plugins.MockMaker;
 import org.mockito.plugins.MockMaker.TypeMockability;
 
 @SuppressWarnings("unchecked")
 public class MockCreationValidator {
 
-    public void validateType(Class<?> classToMock) {
-        TypeMockability typeMockability = MockUtil.typeMockabilityOf(classToMock);
+    public void validateType(Class<?> classToMock, Class<? extends MockMaker> mockMaker) {
+        TypeMockability typeMockability = MockUtil.typeMockabilityOf(classToMock, mockMaker);
         if (!typeMockability.mockable()) {
             throw cannotMockClass(classToMock, typeMockability.nonMockableReason());
         }

diff --git a/src/main/java/org/mockito/internal/util/MockUtil.java b/src/main/java/org/mockito/internal/util/MockUtil.java
@@ -18,22 +18,52 @@
 import org.mockito.plugins.MockMaker.TypeMockability;
 import org.mockito.plugins.MockResolver;
 
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
 import java.util.function.Function;
 
 import static org.mockito.internal.handler.MockHandlerFactory.createMockHandler;
 
 @SuppressWarnings("unchecked")
 public class MockUtil {
 
-    private static final MockMaker mockMaker = Plugins.getMockMaker();
+    private static final MockMaker defaultMockMaker = Plugins.getMockMaker();
+    private static final Map<Class<? extends MockMaker>, MockMaker> mockMakers;
+
+    static {
+        mockMakers = Collections.synchronizedMap(new WeakHashMap<>());
+        mockMakers.put(defaultMockMaker.getClass(), defaultMockMaker);
+    }
 
     private MockUtil() {}
 
+    private static MockMaker getMockMaker(Class<? extends MockMaker> type) {
+        if (type == null) {
+            return defaultMockMaker;
+        }
+        else {
+            return mockMakers.computeIfAbsent(type, t -> {
+                try {
+                    return t.getDeclaredConstructor().newInstance();
+                } catch (Exception e) {
+                    throw new IllegalStateException("Failed to construct MockMaker " + t, e);
+                }
+            });
+        }
+    }
+
     public static TypeMockability typeMockabilityOf(Class<?> type) {
-        return mockMaker.isTypeMockable(type);
+        // TODO: Maybe we should replace all usages of this method with the method below
+        return defaultMockMaker.isTypeMockable(type);
+    }
+
+    public static TypeMockability typeMockabilityOf(Class<?> type, Class<? extends MockMaker> mockMaker) {
+        return getMockMaker(mockMaker).isTypeMockable(type);
     }
 
     public static <T> T createMock(MockCreationSettings<T> settings) {
+        MockMaker mockMaker = getMockMaker(settings.getMockMaker());
         MockHandler mockHandler = createMockHandler(settings);
 
         Object spiedInstance = settings.getSpiedInstance();
@@ -62,17 +92,11 @@ public static void resetMock(Object mock) {
         MockHandler newHandler = createMockHandler(settings);
 
         mock = resolve(mock);
-        mockMaker.resetMock(mock, newHandler, settings);
+        getMockMaker(settings.getMockMaker()).resetMock(mock, newHandler, settings);
     }
 
     public static MockHandler<?> getMockHandler(Object mock) {
-        if (mock == null) {
-            throw new NotAMockException("Argument should be a mock, but is null!");
-        }
-
-        mock = resolve(mock);
-
-        MockHandler handler = mockMaker.getHandler(mock);
+        MockHandler handler = getMockHandler0(mock);
         if (handler != null) {
             return handler;
         } else {
@@ -104,10 +128,24 @@ public static boolean isMock(Object mock) {
         if (mock == null) {
             return false;
         }
+        return getMockHandler0(mock) != null;
+    }
+
+    private static MockHandler<?> getMockHandler0(Object mock) {
+        if (mock == null) {
+            throw new NotAMockException("Argument should be a mock, but is null!");
+        }
 
         mock = resolve(mock);
 
-        return mockMaker.getHandler(mock) != null;
+        for (MockMaker mockMaker : mockMakers.values()) {
+            MockHandler<?> handler = mockMaker.getHandler(mock);
+            if (handler != null) {
+                assert getMockMaker(handler.getMockSettings().getMockMaker()) == mockMaker;
+                return handler;
+            }
+        }
+        return null;
     }
 
     private static Object resolve(Object mock) {
@@ -143,6 +181,7 @@ public static MockCreationSettings getMockSettings(Object mock) {
 
     public static <T> MockMaker.StaticMockControl<T> createStaticMock(
             Class<T> type, MockCreationSettings<T> settings) {
+        MockMaker mockMaker = getMockMaker(settings.getMockMaker());
         MockHandler<T> handler = createMockHandler(settings);
         return mockMaker.createStaticMock(type, settings, handler);
     }
@@ -153,11 +192,13 @@ public static <T> MockMaker.ConstructionMockControl<T> createConstructionMock(
             MockedConstruction.MockInitializer<T> mockInitializer) {
         Function<MockedConstruction.Context, MockHandler<T>> handlerFactory =
                 context -> createMockHandler(settingsFactory.apply(context));
-        return mockMaker.createConstructionMock(
+        return defaultMockMaker.createConstructionMock(
                 type, settingsFactory, handlerFactory, mockInitializer);
     }
 
     public static void clearAllCaches() {
-        mockMaker.clearAllCaches();
+        for (MockMaker mockMaker : mockMakers.values()) {
+            mockMaker.clearAllCaches();
+        }
     }
 }
diff --git a/src/main/java/org/mockito/internal/util/reflection/FieldInitializer.java b/src/main/java/org/mockito/internal/util/reflection/FieldInitializer.java
@@ -259,6 +259,7 @@ public int compare(Constructor<?> constructorA, Constructor<?> constructorB) {
                     private int countMockableParams(Constructor<?> constructor) {
                         int constructorMockableParamsSize = 0;
                         for (Class<?> aClass : constructor.getParameterTypes()) {
+                            // TODO: Should we somehow use the mockMaker from the context?
                             if (MockUtil.typeMockabilityOf(aClass).mockable()) {
                                 constructorMockableParamsSize++;
                             }

diff --git a/src/main/java/org/mockito/mock/MockCreationSettings.java b/src/main/java/org/mockito/mock/MockCreationSettings.java
@@ -12,6 +12,7 @@
 import org.mockito.listeners.InvocationListener;
 import org.mockito.listeners.StubbingLookupListener;
 import org.mockito.listeners.VerificationStartedListener;
+import org.mockito.plugins.MockMaker;
 import org.mockito.quality.Strictness;
 import org.mockito.stubbing.Answer;
 
@@ -135,4 +136,10 @@ public interface MockCreationSettings<T> {
      * @since 4.6.0
      */
     Strictness getStrictness();
+
+    /**
+     * The {@link MockMaker} which shall by used instead of the default. When
+     * the return value is {@code null}, the default shall be used.
+     */
+    Class<? extends MockMaker> getMockMaker();
 }
diff --git a/src/main/java/org/mockito/plugins/MockMaker.java b/src/main/java/org/mockito/plugins/MockMaker.java
@@ -4,6 +4,7 @@
  */
 package org.mockito.plugins;
 
+import org.mockito.MockSettings;
 import org.mockito.MockedConstruction;
 import org.mockito.exceptions.base.MockitoException;
 import org.mockito.invocation.MockHandler;
@@ -45,6 +46,15 @@
  * <p>Note that if several <code>mockito-extensions/org.mockito.plugins.MockMaker</code> files exists in the classpath
  * Mockito will only use the first returned by the standard {@link ClassLoader#getResource} mechanism.
  *
+ * <h3>Using the MockSettings of individual mocks</h3>
+ *
+ * <p>If you want to use your {@code AwesomeMockMaker} only for a specific mock,
+ * you can specify it using {@link MockSettings#mockMaker(Class)}.</p>
+ * <pre>
+ *     Object mock = Mockito.mock(Object.class, Mockito.withSettings()
+ *             .mockMaker(AwesomeMockMaker.class));
+ * </pre>
+ *
  * @see org.mockito.mock.MockCreationSettings
  * @see org.mockito.invocation.MockHandler
  * @since 1.9.5

diff --git a/src/test/java/org/mockito/CustomMockMakerTest.java b/src/test/java/org/mockito/CustomMockMakerTest.java
@@ -0,0 +1,33 @@
+package org.mockito;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.internal.creation.bytebuddy.ByteBuddyMockMaker;
+import org.mockito.internal.creation.bytebuddy.InlineByteBuddyMockMaker;
+
+import java.util.Arrays;
+
+public final class CustomMockMakerTest {
+    @Test
+    public void test_custom_mock_maker() throws Exception {
+        class TestClass {
+            String noop() { return "UNUSED"; }
+            final String finalMethod() { noop(); return "ORIGINAL"; }
+        }
+
+        TestClass inlineMock = Mockito.mock(TestClass.class, Mockito.withSettings()
+            .outerInstance(this)
+            .useConstructor()
+            .mockMaker(InlineByteBuddyMockMaker.class));
+        TestClass subclassMock = Mockito.mock(TestClass.class, Mockito.withSettings()
+            .outerInstance(this)
+            .useConstructor()
+            .mockMaker(ByteBuddyMockMaker.class));
+        for (TestClass mock : Arrays.asList(inlineMock, subclassMock)) {
+            Mockito.when(mock.finalMethod()).thenReturn("MOCKED");
+        }
+
+        Assert.assertEquals("MOCKED", inlineMock.finalMethod());
+        Assert.assertEquals("ORIGINAL", subclassMock.finalMethod());
+    }
+}
diff --git a/src/test/java/org/mockito/internal/util/MockCreationValidatorTest.java b/src/test/java/org/mockito/internal/util/MockCreationValidatorTest.java
@@ -66,7 +66,7 @@ public void should_validation_be_safe_when_nulls_passed() {
     @Test
     public void should_fail_when_type_not_mockable() {
         try {
-            validator.validateType(long.class);
+            validator.validateType(long.class, null);
         } catch (MockitoException ex) {
             assertThat(ex.getMessage()).contains("primitive");
         }