From f19055c0416154c960e9c0974d0a1f76a639f9ea Mon Sep 17 00:00:00 2001
From: Rafael Winterhalter <rafael.wth@gmail.com>
Date: Tue, 2 Jun 2020 14:16:43 +0200
Subject: [PATCH] Add test for post processor.

---
 .../java/net/bytebuddy/asm/AdviceTest.java    | 82 +++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/asm/AdviceTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/asm/AdviceTest.java
index 585b42712e4..d76eddc6d5e 100644
--- a/byte-buddy-dep/src/test/java/net/bytebuddy/asm/AdviceTest.java
+++ b/byte-buddy-dep/src/test/java/net/bytebuddy/asm/AdviceTest.java
@@ -6,8 +6,12 @@
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
 import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.Removal;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.implementation.bytecode.constant.ClassConstant;
+import net.bytebuddy.implementation.bytecode.constant.TextConstant;
+import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess;
 import net.bytebuddy.pool.TypePool;
 import net.bytebuddy.test.packaging.AdviceTestHelper;
 import net.bytebuddy.test.utility.JavaVersionRule;
@@ -1416,6 +1420,60 @@ public void testConstructorNoArgumentBackupAndFrames() throws Exception {
         assertThat(type.getDeclaredConstructor(boolean.class).newInstance(false), notNullValue(Object.class));
     }
 
+    @Test
+    public void testAssigningEnterPostProcessorInline() throws Exception {
+        Class<?> type = new ByteBuddy()
+                .redefine(PostProcessorInlineTest.class)
+                .visit(Advice.withCustomMapping().with(new Advice.PostProcessor.Factory() {
+                    @Override
+                    public Advice.PostProcessor make(final MethodDescription.InDefinedShape advice, boolean exit) {
+                        return new Advice.PostProcessor() {
+                            @Override
+                            public StackManipulation resolve(TypeDescription instrumentedType,
+                                                             MethodDescription instrumentedMethod,
+                                                             Assigner assigner,
+                                                             Advice.ArgumentHandler argumentHandler) {
+                                return new StackManipulation.Compound(
+                                    MethodVariableAccess.of(advice.getReturnType()).loadFrom(argumentHandler.enter()),
+                                    MethodVariableAccess.store(instrumentedMethod.getParameters().get(0))
+                                );
+                            }
+                        };
+                    }
+                }).to(PostProcessorInlineTest.class).on(named(FOO)))
+                .make()
+                .load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
+                .getLoaded();
+        assertThat(type.getDeclaredMethod(FOO, String.class).invoke(type.getDeclaredConstructor().newInstance(), "bar"), is((Object) "foo"));
+    }
+
+    @Test
+    public void testAssigningEnterPostProcessorDelegate() throws Exception {
+        Class<?> type = new ByteBuddy()
+                .redefine(PostProcessorDelegateTest.class)
+                .visit(Advice.withCustomMapping().with(new Advice.PostProcessor.Factory() {
+                    @Override
+                    public Advice.PostProcessor make(final MethodDescription.InDefinedShape advice, boolean exit) {
+                        return new Advice.PostProcessor() {
+                            @Override
+                            public StackManipulation resolve(TypeDescription instrumentedType,
+                                                             MethodDescription instrumentedMethod,
+                                                             Assigner assigner,
+                                                             Advice.ArgumentHandler argumentHandler) {
+                                return new StackManipulation.Compound(
+                                        MethodVariableAccess.of(advice.getReturnType()).loadFrom(argumentHandler.enter()),
+                                        MethodVariableAccess.store(instrumentedMethod.getParameters().get(0))
+                                );
+                            }
+                        };
+                    }
+                }).to(PostProcessorDelegateTest.class).on(named(FOO)))
+                .make()
+                .load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
+                .getLoaded();
+        assertThat(type.getDeclaredMethod(FOO, String.class).invoke(type.getDeclaredConstructor().newInstance(), "bar"), is((Object) "foo"));
+    }
+
     @Test
     public void testPopsValueAfterArrayWrite() throws Exception {
         Class<?> type = new ByteBuddy()
@@ -3416,4 +3474,28 @@ public static void advice(@Advice.AllArguments(readOnly = false, typing = Assign
             String ignored = "" + args;
         }
     }
+
+    public static class PostProcessorInlineTest {
+
+        @Advice.OnMethodEnter
+        static String enter() {
+            return "foo";
+        }
+
+        public String foo(String x) {
+            return x;
+        }
+    }
+
+    public static class PostProcessorDelegateTest {
+
+        @Advice.OnMethodEnter(inline = false)
+        static String enter() {
+            return "foo";
+        }
+
+        public String foo(String x) {
+            return x;
+        }
+    }
 }