Merge pull request mockito#1544 from maxgrabenhorst/fix_premature_cle…

…anup

Fixes mockito#1541: Prevent premature garbage collection of mock objects

src/main/java/org/mockito/internal/stubbing/BaseStubbing.java

@@ -4,17 +4,31 @@
  */
 package org.mockito.internal.stubbing;
 
-import static org.mockito.internal.exceptions.Reporter.notAnException;
-import static org.mockito.internal.progress.ThreadSafeMockingProgress.mockingProgress;
-import static org.objenesis.ObjenesisHelper.newInstance;
-
 import org.mockito.internal.stubbing.answers.CallsRealMethods;
 import org.mockito.internal.stubbing.answers.Returns;
 import org.mockito.internal.stubbing.answers.ThrowsException;
+import org.mockito.stubbing.Answer;
 import org.mockito.stubbing.OngoingStubbing;
 
+import static org.mockito.internal.exceptions.Reporter.notAnException;
+import static org.mockito.internal.progress.ThreadSafeMockingProgress.mockingProgress;
+import static org.objenesis.ObjenesisHelper.newInstance;
+
 public abstract class BaseStubbing<T> implements OngoingStubbing<T> {
 
+
+    // Keep strong ref to mock preventing premature garbage collection when using 'One-liner stubs'. See #1541.
+    private final Object strongMockRef;
+
+    BaseStubbing(Object mock) {
+        this.strongMockRef = mock;
+    }
+
+    @Override
+    public OngoingStubbing<T> then(Answer<?> answer) {
+        return thenAnswer(answer);
+    }
+
     @Override
     public OngoingStubbing<T> thenReturn(T value) {
         return thenAnswer(new Returns(value));
@@ -79,6 +93,12 @@ public OngoingStubbing<T> thenThrow(Class<? extends Throwable> toBeThrown, Class
     public OngoingStubbing<T> thenCallRealMethod() {
         return thenAnswer(new CallsRealMethods());
     }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <M> M getMock() {
+        return (M) this.strongMockRef;
+    }
 }
 
 

src/main/java/org/mockito/internal/stubbing/ConsecutiveStubbing.java

@@ -8,23 +8,17 @@
 import org.mockito.stubbing.OngoingStubbing;
 
 public class ConsecutiveStubbing<T> extends BaseStubbing<T> {
-    private final InvocationContainerImpl invocationContainerImpl;
 
-    public ConsecutiveStubbing(InvocationContainerImpl invocationContainerImpl) {
-        this.invocationContainerImpl = invocationContainerImpl;
+    private final InvocationContainerImpl invocationContainer;
+
+    ConsecutiveStubbing(InvocationContainerImpl invocationContainer) {
+        super(invocationContainer.invokedMock());
+        this.invocationContainer = invocationContainer;
     }
 
     public OngoingStubbing<T> thenAnswer(Answer<?> answer) {
-        invocationContainerImpl.addConsecutiveAnswer(answer);
+        invocationContainer.addConsecutiveAnswer(answer);
         return this;
     }
 
-    public OngoingStubbing<T> then(Answer<?> answer) {
-        return thenAnswer(answer);
-    }
-
-    @SuppressWarnings("unchecked")
-    public <M> M getMock() {
-        return (M) invocationContainerImpl.invokedMock();
-    }
 }

src/main/java/org/mockito/internal/stubbing/OngoingStubbingImpl.java

@@ -19,6 +19,7 @@ public class OngoingStubbingImpl<T> extends BaseStubbing<T> {
     private Strictness strictness;
 
     public OngoingStubbingImpl(InvocationContainerImpl invocationContainer) {
+        super(invocationContainer.invokedMock());
         this.invocationContainer = invocationContainer;
     }
 
@@ -32,22 +33,11 @@ public OngoingStubbing<T> thenAnswer(Answer<?> answer) {
         return new ConsecutiveStubbing<T>(invocationContainer);
     }
 
-    @Override
-    public OngoingStubbing<T> then(Answer<?> answer) {
-        return thenAnswer(answer);
-    }
-
     public List<Invocation> getRegisteredInvocations() {
         //TODO interface for tests
         return invocationContainer.getInvocations();
     }
 
-    @Override
-    @SuppressWarnings("unchecked")
-    public <M> M getMock() {
-        return (M) invocationContainer.invokedMock();
-    }
-
     public void setStrictness(Strictness strictness) {
         this.strictness = strictness;
     }

subprojects/inline/inline.gradle

@@ -9,5 +9,6 @@ dependencies {
 
 tasks.javadoc.enabled = false
 
-//required by the "StressTest.java"
+//required by the "StressTest.java" and "OneLinerStubStressTest.java"
 test.maxHeapSize = "256m"
+retryTest.maxHeapSize = "256m"

subprojects/inline/src/test/java/org/mockitoinline/OneLinerStubStressTest.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) 2018 Mockito contributors
+ * This program is made available under the terms of the MIT License.
+ */
+package org.mockitoinline;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class OneLinerStubStressTest {
+
+    public class OneLinerStubTestClass {
+        public String getStuff() {
+            return "A";
+        }
+    }
+
+    private static String generateLargeString() {
+        final int length = 2000000;
+        final StringBuilder stringBuilder = new StringBuilder(length);
+        for (int i = 0; i <= length; i++) {
+            stringBuilder.append("B");
+        }
+        return stringBuilder.toString();
+    }
+
+    @Test
+    public void call_a_lot_of_mocks_using_one_line_stubbing() {
+        //This requires smaller heap set for the test process, see "inline.gradle"
+        final String returnValue = generateLargeString();
+        for (int i = 0; i < 50000; i++) {
+            // make sure that mock object does not get cleaned up prematurely
+            final OneLinerStubTestClass mock =
+                when(mock(OneLinerStubTestClass.class).getStuff()).thenReturn(returnValue).getMock();
+            assertEquals(returnValue, mock.getStuff());
+        }
+    }
+}