Fix mocks in spies

README.MD

@@ -86,7 +86,7 @@ Simply include the maven dependency (from central maven) to start using @MockInB
 <dependency>
   <groupId>com.teketik</groupId>
   <artifactId>mock-in-bean</artifactId>
-  <version>boot2-v1.5.2</version>
+  <version>boot2-v1.6</version>
   <scope>test</scope>
 </dependency>
 ```

pom.xml

@@ -5,7 +5,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>com.teketik</groupId>
     <artifactId>mock-in-bean</artifactId>
-    <version>boot2-v1.5.2</version>
+    <version>boot2-v1.6</version>
     <name>Mock in Bean</name>
     <description>Surgically Inject Mockito Mock/Spy in Spring Beans</description>
     <url>https://github.com/antoinemeyer/mock-in-bean</url>

src/main/java/com/teketik/test/mockinbean/MockInBeanTestExecutionListener.java

@@ -15,10 +15,12 @@
 import java.lang.reflect.Type;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -94,22 +96,44 @@ public void beforeTestMethod(TestContext testContext) throws Exception {
         final TestContext applicableTestContext = ROOT_TEST_CONTEXT_TRACKER
                 .get(resolveTestClass(testContext.getTestClass()));
         final Map<Definition, Object> mockOrSpys = new HashMap<>();
-        ((LinkedList<FieldState>) applicableTestContext.getAttribute(ORIGINAL_VALUES_ATTRIBUTE_NAME))
+        final LinkedList<FieldState> fieldStates = (LinkedList<FieldState>) applicableTestContext.getAttribute(ORIGINAL_VALUES_ATTRIBUTE_NAME);
+        final Map<Object, Object> spyTracker = new IdentityHashMap<>();
+        //First loop to setup all the mocks and spies
+        fieldStates
             .forEach(fieldState -> {
                 Object mockOrSpy = mockOrSpys.get(fieldState.definition);
                 if (mockOrSpy == null) {
                     mockOrSpy = fieldState.definition.create(fieldState.originalValue);
                     mockOrSpys.put(fieldState.definition, mockOrSpy);
+                    if (fieldState.definition instanceof SpyDefinition) {
+                        spyTracker.put(fieldState.originalValue, mockOrSpy);
+                    }
                 }
-                ReflectionUtils.setField(
-                    fieldState.field,
-                    fieldState.resolveTarget(applicableTestContext),
-                    mockOrSpy
-                );
             });
+        //Second loop to process the injections (handling mocks in spies)
+        fieldStates
+            .forEach(fieldState -> {
+                final Object mockOrSpy = mockOrSpys.get(fieldState.definition);
+                final Object bean = fieldState.resolveTarget(applicableTestContext);
+                //inject in original bean
+                inject(fieldState.field, bean, mockOrSpy);
+                //if the target bean has been spied on, need to push into this spy as well (to allow mock in spies)
+                Optional.ofNullable(spyTracker.get(bean))
+                    .ifPresent(spy ->inject(fieldState.field, spy, mockOrSpy));
+
+            });
+
         super.beforeTestMethod(testContext);
     }
 
+    private void inject(Field field, Object inObject, Object toInject) {
+        ReflectionUtils.setField(
+            field,
+            inObject,
+            toInject
+        );
+    }
+
     /*
      * Iterate over all the definitions and put back the original values in the beans
      */

src/test/java/com/teketik/test/mockinbean/test/MockInSpyTest.java

@@ -0,0 +1,76 @@
+package com.teketik.test.mockinbean.test;
+
+import com.teketik.test.mockinbean.MockInBean;
+import com.teketik.test.mockinbean.SpyInBean;
+import com.teketik.test.mockinbean.test.MockInSpyTest.Config.TestComponent1Wrapper;
+import com.teketik.test.mockinbean.test.components.MockableComponent1;
+import com.teketik.test.mockinbean.test.components.TestComponent1;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import org.springframework.stereotype.Component;
+
+import javax.annotation.Resource;
+
+class MockInSpyTest extends BaseTest {
+
+    @org.springframework.boot.test.context.TestConfiguration
+    static class Config {
+
+        @Component
+        class TestComponent1Wrapper {
+
+            @Resource
+            private TestComponent1 testComponent1;
+
+            void doWith1() {
+                testComponent1.doWith1();
+            }
+
+            void doWith2() {
+                testComponent1.doWith2();
+            }
+
+        }
+    }
+
+    @Resource
+    private TestComponent1Wrapper testComponent1Wrapper;
+
+    @SpyInBean(TestComponent1Wrapper.class)
+    private TestComponent1 testComponent1;
+
+    @MockInBean(TestComponent1.class)
+    private MockableComponent1 mockableComponent1;
+
+    @Test
+    public void testMocked1() {
+        Mockito.doNothing().when(mockableComponent1).doSomething();
+        testComponent1Wrapper.doWith1();
+        Mockito.verify(mockableComponent1).doSomething();
+    }
+
+    @Test
+    public void testMocked2() {
+        final RuntimeException runtimeException = new RuntimeException();
+        Mockito.doThrow(runtimeException).when(mockableComponent1).doSomething();
+        try {
+            testComponent1Wrapper.doWith1();
+            Assertions.fail();
+        } catch (Exception e) {
+            Assertions.assertSame(runtimeException, e);
+        }
+    }
+
+    @Test
+    public void testNotMocked() {
+        try {
+            testComponent1Wrapper.doWith2();
+            Assertions.fail();
+        } catch (Exception e) {
+            Assertions.assertTrue(e instanceof UnsupportedOperationException);
+        }
+    }
+
+}

src/test/java/com/teketik/test/mockinbean/test/components/TestComponentBase.java

@@ -10,4 +10,12 @@ public abstract class TestComponentBase {
     @Autowired
     private MockableComponent2 mockableComponent2;
 
+    public void doWith1() {
+        mockableComponent1.doSomething();
+    }
+
+    public void doWith2() {
+        mockableComponent2.doSomething();
+    }
+
 }