From f4cc94d2b1526dc53546233944612c38897fdb65 Mon Sep 17 00:00:00 2001
From: "ah.jo" <ah.jo@navercorp.com>
Date: Tue, 31 Dec 2024 10:51:32 +0900
Subject: [PATCH] Fix not applying registered child manipulations if parent
 uses thenApply

---
 .../fixturemonkey/tests/kotlin/KotlinTest.kt  | 20 ++++++++
 .../fixturemonkey/tree/ObjectNode.java        | 51 ++++++++++++++++---
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/fixture-monkey-tests/kotlin-tests/src/test/kotlin/com/navercorp/fixturemonkey/tests/kotlin/KotlinTest.kt b/fixture-monkey-tests/kotlin-tests/src/test/kotlin/com/navercorp/fixturemonkey/tests/kotlin/KotlinTest.kt
index bc8206de88..7da2a9dc0d 100644
--- a/fixture-monkey-tests/kotlin-tests/src/test/kotlin/com/navercorp/fixturemonkey/tests/kotlin/KotlinTest.kt
+++ b/fixture-monkey-tests/kotlin-tests/src/test/kotlin/com/navercorp/fixturemonkey/tests/kotlin/KotlinTest.kt
@@ -1306,6 +1306,26 @@ class KotlinTest {
         then(actual).isNotNull
     }
 
+    @Test
+    fun parentRegisterThenApplyAndSizeReturnsChildRegister() {
+        data class ChildObject(val value: String)
+
+        data class ParentObject(val values: List<ChildObject>)
+
+        val sut = FixtureMonkey.builder()
+            .register { it.giveMeKotlinBuilder<ParentObject>().thenApply { _, _ -> } }
+            .register { it.giveMeKotlinBuilder<ChildObject>().set(ChildObject::value, "1") }
+            .plugin(KotlinPlugin())
+            .build()
+
+        val actual = sut.giveMeKotlinBuilder<ParentObject>()
+            .size(ParentObject::values, 10)
+            .sample()
+            .values
+
+        then(actual).allMatch { it.value == "1" }
+    }
+
     companion object {
         private val SUT: FixtureMonkey = FixtureMonkey.builder()
             .plugin(KotlinPlugin())
diff --git a/fixture-monkey/src/main/java/com/navercorp/fixturemonkey/tree/ObjectNode.java b/fixture-monkey/src/main/java/com/navercorp/fixturemonkey/tree/ObjectNode.java
index 7214886ae9..9b722ddb03 100644
--- a/fixture-monkey/src/main/java/com/navercorp/fixturemonkey/tree/ObjectNode.java
+++ b/fixture-monkey/src/main/java/com/navercorp/fixturemonkey/tree/ObjectNode.java
@@ -19,8 +19,12 @@
 package com.navercorp.fixturemonkey.tree;
 
 import static com.navercorp.fixturemonkey.api.type.Types.nullSafe;
+import static java.util.stream.Collectors.toMap;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import javax.annotation.Nullable;
@@ -29,6 +33,7 @@
 import org.apiguardian.api.API.Status;
 
 import com.navercorp.fixturemonkey.api.generator.ArbitraryProperty;
+import com.navercorp.fixturemonkey.api.generator.ObjectProperty;
 import com.navercorp.fixturemonkey.api.lazy.LazyArbitrary;
 import com.navercorp.fixturemonkey.api.property.Property;
 import com.navercorp.fixturemonkey.api.property.PropertyPath;
@@ -83,9 +88,11 @@ public boolean expand() {
 	public void forceExpand() {
 		this.traverseNode.forceExpand();
 		this.setChildren(
-			nullSafe(this.traverseNode.getChildren()).asList().stream()
-				.map(it -> new ObjectNode(it, generateFixtureContext.newChildNodeContext()))
-				.collect(Collectors.toList())
+			this.mergeWithNewChildren(
+				nullSafe(this.traverseNode.getChildren()).asList().stream()
+					.map(it -> new ObjectNode(it, generateFixtureContext.newChildNodeContext()))
+					.collect(Collectors.toList())
+			)
 		);
 	}
 
@@ -93,9 +100,11 @@ public void forceExpand() {
 	public void forceExpand(TypeDefinition typeDefinition) {
 		this.traverseNode.forceExpand(typeDefinition);
 		this.setChildren(
-			nullSafe(this.traverseNode.getChildren()).asList().stream()
-				.map(it -> new ObjectNode(it, generateFixtureContext.newChildNodeContext()))
-				.collect(Collectors.toList())
+			this.mergeWithNewChildren(
+				nullSafe(this.traverseNode.getChildren()).asList().stream()
+					.map(it -> new ObjectNode(it, generateFixtureContext.newChildNodeContext()))
+					.collect(Collectors.toList())
+			)
 		);
 	}
 
@@ -191,4 +200,34 @@ public TreeNodeManipulator getAppliedTreeNodeManipulator() {
 	public GenerateFixtureContext getObjectNodeContext() {
 		return generateFixtureContext;
 	}
+
+	private List<ObjectNode> mergeWithNewChildren(List<ObjectNode> newChildren) {
+		if (this.children == null) {
+			return newChildren;
+		}
+
+		boolean shrinkChildNodes = this.children.size() > newChildren.size();
+		if (shrinkChildNodes) {
+			return this.children.subList(0, newChildren.size());
+		}
+
+		boolean expandChildNodes = this.children.size() < newChildren.size();
+		if (expandChildNodes) {
+			Map<ObjectProperty, ObjectNode> existingNodesByObjectProperty = this.children.stream()
+				.collect(toMap(it -> it.getMetadata().getTreeProperty().getObjectProperty(), Function.identity()));
+
+			List<ObjectNode> concatNewChildren = new ArrayList<>();
+			for (ObjectNode newChild : newChildren) {
+				ObjectNode existingNode =
+					existingNodesByObjectProperty.get(newChild.getMetadata().getTreeProperty().getObjectProperty());
+				if (existingNode != null) {
+					concatNewChildren.add(existingNode);
+				} else {
+					concatNewChildren.add(newChild);
+				}
+			}
+			return concatNewChildren;
+		}
+		return this.children;
+	}
 }