maintain original aspect functionality

antoinemeyer · Nov 13, 2024 · 4260471 · 4260471
1 parent 5735f80
commit 4260471
Show file tree

Hide file tree

Showing 5 changed files with 110 additions and 40 deletions.
diff --git a/src/main/java/com/teketik/test/mockinbean/BeanFieldState.java b/src/main/java/com/teketik/test/mockinbean/BeanFieldState.java
@@ -1,34 +1,34 @@
 package com.teketik.test.mockinbean;
 
 import org.springframework.test.context.TestContext;
+import org.springframework.util.ReflectionUtils;
 
 import java.lang.reflect.Field;
 
 class BeanFieldState extends FieldState {
 
-    private Object bean;
+    final Object bean;
 
-    private Object originalValue;
+    final Object originalValue;
 
-    private Object mockableValue;
-
-    public BeanFieldState(Object bean, Field field, Object originalValue, Object mockableValue, Definition definition) {
+    public BeanFieldState(Object bean, Field field, Object originalValue, Definition definition) {
         super(field, definition);
         this.bean = bean;
         this.originalValue = originalValue;
-        this.mockableValue = mockableValue;
     }
 
     @Override
     public Object resolveTarget(TestContext testContext) {
         return bean;
     }
 
-    public Object getMockableValue() {
-        return mockableValue;
+    public void rollback(TestContext testContext) {
+        final Object target = resolveTarget(testContext);
+        ReflectionUtils.setField(field, target, originalValue);
     }
 
-    public Object getOriginalValue() {
-        return originalValue;
+    public Object createMockOrSpy() {
+        return definition.create(originalValue);
     }
+
 }
diff --git a/src/main/java/com/teketik/test/mockinbean/BeanUtils.java b/src/main/java/com/teketik/test/mockinbean/BeanUtils.java
@@ -1,5 +1,7 @@
 package com.teketik.test.mockinbean;
 
+import org.springframework.aop.TargetSource;
+import org.springframework.aop.framework.Advised;
 import org.springframework.aop.framework.AopProxyUtils;
 import org.springframework.aop.support.AopUtils;
 import org.springframework.context.ApplicationContext;
@@ -44,8 +46,8 @@ static <T> T findBean(Class<T> type, @Nullable String name, ApplicationContext a
                 .findFirst()
                 .orElseThrow(() -> new IllegalArgumentException("No beans of type " + type + " and name " + name));
         }
-        return AopUtils.isAopProxy(beanOrProxy) 
-            ? (T) AopProxyUtils.getSingletonTarget(beanOrProxy) 
+        return AopUtils.isAopProxy(beanOrProxy)
+            ? (T) AopProxyUtils.getSingletonTarget(beanOrProxy)
             : beanOrProxy;
     }
 
@@ -96,4 +98,27 @@ static Field findField(Class<?> clazz, @Nullable String name, Class<?> type) {
         return null;
     }
 
+    static @Nullable TargetSource getProxyTarget(Object candidate) {
+        try {
+            while (AopUtils.isAopProxy(candidate) && candidate instanceof Advised) {
+                Advised advised = (Advised) candidate;
+                TargetSource targetSource = advised.getTargetSource();
+
+                if (targetSource.isStatic()) {
+                    Object target = targetSource.getTarget();
+
+                    if (target == null || !AopUtils.isAopProxy(target)) {
+                        return targetSource;
+                    }
+                    candidate = target;
+                } else {
+                    return null;
+                }
+            }
+        } catch (Throwable ex) {
+            throw new IllegalStateException("Failed to unwrap proxied object", ex);
+        }
+        return null;
+    }
+
 }
diff --git a/src/main/java/com/teketik/test/mockinbean/MockInBeanTestExecutionListener.java b/src/main/java/com/teketik/test/mockinbean/MockInBeanTestExecutionListener.java
@@ -4,8 +4,6 @@
 import org.mockito.Mock;
 import org.mockito.Spy;
 import org.springframework.aop.TargetSource;
-import org.springframework.aop.framework.Advised;
-import org.springframework.aop.support.AopUtils;
 import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.test.context.TestContext;
 import org.springframework.test.context.TestExecutionListener;
@@ -62,12 +60,18 @@ public void beforeTestClass(TestContext testContext) throws Exception {
             for (InBeanDefinition inBeanDefinition : definitionToInbeans.getValue()) {
                 final Object inBean = BeanUtils.findBean(inBeanDefinition.clazz, inBeanDefinition.name, testContext.getApplicationContext());
                 beanField = BeanUtils.findField(inBean.getClass(), definition.getName(), mockOrSpyType);
+                Assert.notNull(beanField, "Cannot find any field for definition:" + definitionToInbeans.getKey());
                 beanField.setAccessible(true);
                 final Object beanFieldValue = ReflectionUtils.getField(beanField, inBean);
-                originalValues.add(new BeanFieldState(inBean, beanField, beanFieldValue,
-                    determineMockableValue(beanFieldValue), definition));
+                final TargetSource proxyTarget = BeanUtils.getProxyTarget(beanFieldValue);
+                BeanFieldState beanFieldState;
+                if (proxyTarget != null) {
+                    beanFieldState = new ProxiedBeanFieldState(inBean, beanField, beanFieldValue, proxyTarget, definition);
+                } else {
+                    beanFieldState = new BeanFieldState(inBean, beanField, beanFieldValue, definition);
+                }
+                originalValues.add(beanFieldState);
             }
-            Assert.notNull(beanField, "Cannot find any field for definition:" + definitionToInbeans.getKey());
             Assert.isTrue(visitedFields.add(beanField), beanField + " can only be mapped once, as a mock or a spy, not both!");
             final Field testField = ReflectionUtils.findField(targetTestClass, definition.getName(), mockOrSpyType);
             testField.setAccessible(true);
@@ -77,21 +81,6 @@ public void beforeTestClass(TestContext testContext) throws Exception {
         super.beforeTestClass(testContext);
     }
 
-    private Object determineMockableValue(Object candidate) {
-        try {
-            while (candidate != null
-                    && AopUtils.isAopProxy(candidate)
-                    && candidate instanceof Advised) {
-                final Advised advised = (Advised) candidate;
-                final TargetSource targetSource = advised.getTargetSource();
-                candidate = targetSource.getTarget();
-            }
-            return candidate;
-        } catch (Exception e) {
-            throw new RuntimeException("Target cannot be resolved", e);
-        }
-    }
-
     /*
      * Iterate over all the definitions and create a corresponding mock/spy that is injected in the beans and the test class
      */
@@ -110,10 +99,10 @@ public void beforeTestMethod(TestContext testContext) throws Exception {
             .forEach(fieldState -> {
                 Object mockOrSpy = mockOrSpys.get(fieldState.definition);
                 if (mockOrSpy == null) {
-                    mockOrSpy = fieldState.definition.create(fieldState.getMockableValue());
+                    mockOrSpy = fieldState.createMockOrSpy();
                     mockOrSpys.put(fieldState.definition, mockOrSpy);
                     if (fieldState.definition instanceof SpyDefinition) {
-                        spyTracker.put(fieldState.getMockableValue(), mockOrSpy);
+                        spyTracker.put(fieldState.originalValue, mockOrSpy);
                     }
                 }
             });
@@ -152,8 +141,7 @@ public void afterTestClass(TestContext testContext) throws Exception {
             .stream()
             .filter(BeanFieldState.class::isInstance)
             .map(BeanFieldState.class::cast)
-            .forEach(fieldValue ->
-                ReflectionUtils.setField(fieldValue.field, fieldValue.resolveTarget(testContext), fieldValue.getOriginalValue()));
+            .forEach(fieldState -> fieldState.rollback(testContext));
         ROOT_TEST_CONTEXT_TRACKER.remove(testContext.getTestClass());
         super.afterTestClass(testContext);
     }

diff --git a/src/main/java/com/teketik/test/mockinbean/ProxiedBeanFieldState.java b/src/main/java/com/teketik/test/mockinbean/ProxiedBeanFieldState.java
@@ -0,0 +1,42 @@
+package com.teketik.test.mockinbean;
+
+import org.springframework.aop.TargetSource;
+import org.springframework.test.context.TestContext;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import java.lang.reflect.Field;
+
+/**
+ * Special kind of {@link BeanFieldState} handling proxied beans (like aspects).<br>
+ * The mock is not injected into the <code>field</code> but into the <code>target</code> of its {@link TargetSource}.
+ * @author Antoine Meyer
+ */
+class ProxiedBeanFieldState extends BeanFieldState {
+
+    private static void setTargetSourceValue(TargetSource targetSource, Object value) {
+        ReflectionTestUtils.setField(targetSource, "target", value);
+    }
+
+    final TargetSource proxyTargetSource;
+
+    final Object proxyTargetOriginalValue;
+
+    public ProxiedBeanFieldState(Object inBean, Field beanField, Object beanFieldValue, TargetSource proxyTargetSource, Definition definition) throws Exception {
+        super(inBean, beanField, beanFieldValue, definition);
+        this.proxyTargetSource = proxyTargetSource;
+        this.proxyTargetOriginalValue = proxyTargetSource.getTarget();
+    }
+
+    @Override
+    public void rollback(TestContext testContext) {
+        setTargetSourceValue(proxyTargetSource, proxyTargetOriginalValue);
+    }
+
+    @Override
+    public Object createMockOrSpy() {
+        Object applicableMockOrSpy = definition.create(proxyTargetOriginalValue);
+        setTargetSourceValue(proxyTargetSource, applicableMockOrSpy);
+        return originalValue; //the 'mock or spy' to operate for proxied beans are the actual proxy
+    }
+
+}
diff --git a/src/test/java/com/teketik/test/mockinbean/test/VerifyAdvisedSpyInBeanTest.java b/src/test/java/com/teketik/test/mockinbean/test/VerifyAdvisedSpyInBeanTest.java
@@ -3,6 +3,7 @@
 import static org.mockito.Mockito.verify;
 
 import com.teketik.test.mockinbean.SpyInBean;
+import com.teketik.test.mockinbean.test.VerifyAdvisedSpyInBeanTest.Config.AnAspect;
 import com.teketik.test.mockinbean.test.VerifyAdvisedSpyInBeanTest.Config.LoggingService;
 import com.teketik.test.mockinbean.test.VerifyAdvisedSpyInBeanTest.Config.ProviderService;
 
@@ -22,9 +23,10 @@
 import org.springframework.test.context.TestExecutionListeners.MergeMode;
 import org.springframework.test.util.ReflectionTestUtils;
 
+import java.util.concurrent.atomic.AtomicInteger;
+
 /**
- * Covering test case from
- * https://github.com/inkassso/mock-in-bean-issue-23/blob/master/src/test/java/com/github/inkassso/mockinbean/issue23/service/BrokenLoggingServiceTest1_SpyInBean.java
+ * Covering test case from https://github.com/inkassso/mock-in-bean-issue-23/blob/master/src/test/java/com/github/inkassso/mockinbean/issue23/service/BrokenLoggingServiceTest1_SpyInBean.java
  */
 @TestExecutionListeners(value = {VerifyAdvisedSpyInBeanTest.class}, mergeMode = MergeMode.MERGE_WITH_DEFAULTS)
 @SpringBootTest
@@ -36,8 +38,13 @@ static class Config {
         @Aspect
         @Component
         public class AnAspect {
+
+            private final AtomicInteger invocationCounter = new AtomicInteger();
+
             @Before("execution(* com.teketik.test.mockinbean.test.VerifyAdvisedSpyInBeanTest.Config.ProviderService.provideValue())")
-            public void logBeforeMethodExecution() {}
+            public void run() {
+                invocationCounter.incrementAndGet();
+            }
         }
 
         @Service
@@ -74,8 +81,16 @@ void testLogCurrentValue() {
     @Override
     public void afterTestClass(TestContext testContext) throws Exception {
         final ApplicationContext applicationContext = testContext.getApplicationContext();
+
+        //ensure context clean
         final Object loggingServiceBean = applicationContext.getBean(LoggingService.class);
-        Assertions.assertSame(applicationContext.getBean(ProviderService.class), ReflectionTestUtils.getField(loggingServiceBean, "providerService"));
+        final Object providerServiceInBean = ReflectionTestUtils.getField(loggingServiceBean, "providerService");
+        Assertions.assertFalse(TestUtils.isMockOrSpy(providerServiceInBean));
+        Assertions.assertSame(applicationContext.getBean(ProviderService.class), providerServiceInBean);
+
+        //ensure aspect invoked (from log and verify)
+        final AnAspect anAspect = applicationContext.getBean(AnAspect.class);
+        Assertions.assertEquals(2, anAspect.invocationCounter.get());
     }
 
     @Override