Merge pull request #30 from mockito/fix-28

save values of InjectMock fields before test
mockito · Dec 13, 2020 · f5f0707 · f5f0707
2 parents a204794 + dd17a7f
commit f5f0707
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 12 deletions.
diff --git a/src/main/java/org/mockito/testng/MockitoTestNGListener.java b/src/main/java/org/mockito/testng/MockitoTestNGListener.java
@@ -6,13 +6,16 @@
 
 import java.lang.annotation.Annotation;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
-import java.util.WeakHashMap;
 import java.util.stream.Stream;
 
+import org.mockito.InjectMocks;
 import org.mockito.Mockito;
 import org.mockito.MockitoSession;
+import org.mockito.internal.util.reflection.Fields;
+import org.mockito.internal.util.reflection.InstanceField;
 import org.mockito.quality.Strictness;
 import org.testng.IInvokedMethod;
 import org.testng.IInvokedMethodListener;
@@ -77,16 +80,28 @@
  */
 public class MockitoTestNGListener implements IInvokedMethodListener {
 
-    private final Map<Object, MockitoSession> sessions = new WeakHashMap<>();
+    private final Map<Object, MockitoSession> sessions = new HashMap<>();
+    private final Map<Object, Map<InstanceField, Object>> injectMocksFieldsValues = new HashMap<>();
 
     @Override
     public void beforeInvocation(IInvokedMethod method, ITestResult testResult) {
         if (shouldBeRunBeforeInvocation(method, testResult)) {
+
+            // save value of all InjectMocks fields
+            // in order to restore state before next tests
+            // https://github.com/mockito/mockito-testng/issues/28
+            injectMocksFieldsValues.computeIfAbsent(testResult.getInstance(), testInstance ->
+                    Fields.allDeclaredFieldsOf(testInstance).instanceFields()
+                            .stream()
+                            .filter(field -> field.isAnnotatedBy(InjectMocks.class))
+                            .collect(HashMap::new, (m, v) -> m.put(v, v.read()), HashMap::putAll));
+
             sessions.computeIfAbsent(testResult.getInstance(), testInstance -> {
 
                         Strictness strictness = findAnnotation(testResult, MockitoSettings.class)
                                 .map(MockitoSettings::strictness).orElse(Strictness.STRICT_STUBS);
 
+                        // start MockitoSession
                         return Mockito.mockitoSession()
                                 .initMocks(testInstance)
                                 .strictness(strictness)
@@ -99,8 +114,13 @@ public void beforeInvocation(IInvokedMethod method, ITestResult testResult) {
     @Override
     public void afterInvocation(IInvokedMethod method, ITestResult testResult) {
         if (shouldBeRunAfterInvocation(method, testResult)) {
-            Optional.ofNullable(sessions.remove(testResult.getInstance()))
-                    .ifPresent(mockitoSession -> mockitoSession.finishMocking(testResult.getThrowable()));
+            try {
+                Optional.ofNullable(sessions.remove(testResult.getInstance()))
+                        .ifPresent(mockitoSession -> mockitoSession.finishMocking(testResult.getThrowable()));
+            } finally {
+                Optional.ofNullable(injectMocksFieldsValues.remove(testResult.getInstance()))
+                        .ifPresent(fieldsValues -> fieldsValues.forEach(InstanceField::set));
+            }
         }
     }
 

diff --git a/src/test/java/org/mockitousage/testng/InjectMocksWithConstructorAndFinalTest.java b/src/test/java/org/mockitousage/testng/InjectMocksWithConstructorAndFinalTest.java
@@ -14,7 +14,6 @@
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.testng.MockitoTestNGListener;
-import org.testng.annotations.AfterMethod;
 import org.testng.annotations.Listeners;
 import org.testng.annotations.Test;
 
@@ -44,13 +43,8 @@ int callClient() {
     private Service service;
 
     private int lastClientHash = 0;
+    private int lastServiceHash = 0;
 
-    @AfterMethod
-    void cleanup() {
-        // final field will be not assigned for next test
-        // for new object constructor will be used
-        service = null;
-    }
 
     @Test(invocationCount = 4)
     void inject_mock_should_be_refreshed() {
@@ -61,6 +55,8 @@ void inject_mock_should_be_refreshed() {
 
         // we have new mock
         assertThat(lastClientHash).isNotEqualTo(client.hashCode());
+        // we have new InjectMocks filed value
+        assertThat(lastServiceHash).isNotEqualTo(service.hashCode());
 
         // clear mock
         assertThat(service.callClient()).isZero();
@@ -74,7 +70,8 @@ void inject_mock_should_be_refreshed() {
 
         verify(client, times(2)).aMethod();
 
-        // remember last mock hash
+        // remember last mock and InjectMocks field hash
         lastClientHash = client.hashCode();
+        lastServiceHash = service.hashCode();
     }
 }