From 7cfac7c060390decec58b0ea68a29d7ac6378a5c Mon Sep 17 00:00:00 2001
From: Zegveld <41897697+Zegveld@users.noreply.github.com>
Date: Sun, 2 Oct 2022 09:34:03 +0200
Subject: [PATCH] #2955 Fix `@AfterMapping` with return type not called for
 update mappings

---
 .../model/LifecycleMethodResolver.java        |  2 +-
 .../CallbacksWithReturnValuesTest.java        | 46 +++++++++++++++----
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/processor/src/main/java/org/mapstruct/ap/internal/model/LifecycleMethodResolver.java b/processor/src/main/java/org/mapstruct/ap/internal/model/LifecycleMethodResolver.java
index 87527e87f..8b44dee25 100644
--- a/processor/src/main/java/org/mapstruct/ap/internal/model/LifecycleMethodResolver.java
+++ b/processor/src/main/java/org/mapstruct/ap/internal/model/LifecycleMethodResolver.java
@@ -141,7 +141,7 @@ public final class LifecycleMethodResolver {
             callbackMethods,
             Collections.emptyList(),
             targetType,
-            method.getReturnType(),
+            method.getResultType(),
             SelectionCriteria.forLifecycleMethods( selectionParameters ) );
 
         return toLifecycleCallbackMethodRefs(
diff --git a/processor/src/test/java/org/mapstruct/ap/test/callbacks/returning/CallbacksWithReturnValuesTest.java b/processor/src/test/java/org/mapstruct/ap/test/callbacks/returning/CallbacksWithReturnValuesTest.java
index 8c391bbfe..7d6eaacb7 100644
--- a/processor/src/test/java/org/mapstruct/ap/test/callbacks/returning/CallbacksWithReturnValuesTest.java
+++ b/processor/src/test/java/org/mapstruct/ap/test/callbacks/returning/CallbacksWithReturnValuesTest.java
@@ -8,6 +8,7 @@ package org.mapstruct.ap.test.callbacks.returning;
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicReference;
 
+import org.junit.jupiter.api.AfterEach;
 import org.mapstruct.ap.test.callbacks.returning.NodeMapperContext.ContextListener;
 import org.mapstruct.ap.testutil.IssueKey;
 import org.mapstruct.ap.testutil.ProcessorTest;
@@ -25,23 +26,29 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
 @WithClasses( { Attribute.class, AttributeDto.class, Node.class, NodeDto.class, NodeMapperDefault.class,
     NodeMapperWithContext.class, NodeMapperContext.class, Number.class, NumberMapperDefault.class,
     NumberMapperContext.class, NumberMapperWithContext.class } )
-public class CallbacksWithReturnValuesTest {
+class CallbacksWithReturnValuesTest {
+    @AfterEach
+    void cleanup() {
+        NumberMapperContext.clearCache();
+        NumberMapperContext.clearVisited();
+    }
+
     @ProcessorTest
-    public void mappingWithDefaultHandlingRaisesStackOverflowError() {
+    void mappingWithDefaultHandlingRaisesStackOverflowError() {
         Node root = buildNodes();
         assertThatThrownBy( () -> NodeMapperDefault.INSTANCE.nodeToNodeDto( root ) )
             .isInstanceOf( StackOverflowError.class );
     }
 
     @ProcessorTest
-    public void updatingWithDefaultHandlingRaisesStackOverflowError() {
+    void updatingWithDefaultHandlingRaisesStackOverflowError() {
         Node root = buildNodes();
         assertThatThrownBy( () -> NodeMapperDefault.INSTANCE.nodeToNodeDto( root, new NodeDto() ) )
             .isInstanceOf( StackOverflowError.class );
     }
 
     @ProcessorTest
-    public void mappingWithContextCorrectlyResolvesCycles() {
+    void mappingWithContextCorrectlyResolvesCycles() {
         final AtomicReference<Integer> contextLevel = new AtomicReference<>( null );
         ContextListener contextListener = new ContextListener() {
             @Override
@@ -75,28 +82,49 @@ public class CallbacksWithReturnValuesTest {
     }
 
     @ProcessorTest
-    public void numberMappingWithoutContextDoesNotUseCache() {
+    void numberMappingWithoutContextDoesNotUseCache() {
         Number n1 = NumberMapperDefault.INSTANCE.integerToNumber( 2342 );
         Number n2 = NumberMapperDefault.INSTANCE.integerToNumber( 2342 );
+
         assertThat( n1 ).isEqualTo( n2 );
         assertThat( n1 ).isNotSameAs( n2 );
     }
 
     @ProcessorTest
-    public void numberMappingWithContextUsesCache() {
+    void numberMappingWithContextUsesCache() {
         NumberMapperContext.putCache( new Number( 2342 ) );
         Number n1 = NumberMapperWithContext.INSTANCE.integerToNumber( 2342 );
         Number n2 = NumberMapperWithContext.INSTANCE.integerToNumber( 2342 );
+
         assertThat( n1 ).isEqualTo( n2 );
         assertThat( n1 ).isSameAs( n2 );
-        NumberMapperContext.clearCache();
     }
 
     @ProcessorTest
-    public void numberMappingWithContextCallsVisitNumber() {
+    void numberMappingWithContextCallsVisitNumber() {
         Number n1 = NumberMapperWithContext.INSTANCE.integerToNumber( 1234 );
         Number n2 = NumberMapperWithContext.INSTANCE.integerToNumber( 5678 );
+
         assertThat( NumberMapperContext.getVisited() ).isEqualTo( Arrays.asList( n1, n2 ) );
-        NumberMapperContext.clearVisited();
+    }
+
+    @ProcessorTest
+    @IssueKey( "2955" )
+    void numberUpdateMappingWithContextUsesCacheAndThereforeDoesNotVisitNumber() {
+        Number target = new Number();
+        Number expectedReturn = new Number( 2342 );
+        NumberMapperContext.putCache( expectedReturn );
+        NumberMapperWithContext.INSTANCE.integerToNumber( 2342, target );
+
+        assertThat( NumberMapperContext.getVisited() ).isEmpty();
+    }
+
+    @ProcessorTest
+    @IssueKey( "2955" )
+    void numberUpdateMappingWithContextCallsVisitNumber() {
+        Number target = new Number();
+        NumberMapperWithContext.INSTANCE.integerToNumber( 2342, target );
+
+        assertThat( NumberMapperContext.getVisited() ).contains( target );
     }
 }