From 3da533d245356aec32f36aaf1cd1f4ca470140ca Mon Sep 17 00:00:00 2001
From: rstoyanchev <rossen.stoyanchev@broadcom.com>
Date: Wed, 20 Mar 2024 20:14:37 +0000
Subject: [PATCH] Update EntityMapping test to use batch mapping

---
 .../federation/EntityMappingInvocationTests.java | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java
index c8bf9fcd..e008fbde 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/data/federation/EntityMappingInvocationTests.java
@@ -23,6 +23,7 @@ import graphql.GraphQLError;
 import graphql.GraphqlErrorBuilder;
 import graphql.schema.DataFetchingEnvironment;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
@@ -38,9 +39,11 @@ import org.springframework.graphql.ResponseHelper;
 import org.springframework.graphql.TestExecutionGraphQlService;
 import org.springframework.graphql.TestExecutionRequest;
 import org.springframework.graphql.data.method.annotation.Argument;
+import org.springframework.graphql.data.method.annotation.BatchMapping;
 import org.springframework.graphql.data.method.annotation.GraphQlExceptionHandler;
-import org.springframework.graphql.data.method.annotation.SchemaMapping;
 import org.springframework.graphql.data.method.annotation.support.AnnotatedControllerConfigurer;
+import org.springframework.graphql.execution.BatchLoaderRegistry;
+import org.springframework.graphql.execution.DefaultBatchLoaderRegistry;
 import org.springframework.graphql.execution.ErrorType;
 import org.springframework.lang.Nullable;
 import org.springframework.stereotype.Controller;
@@ -130,8 +133,11 @@ public class EntityMappingInvocationTests {
 	}
 
 	private TestExecutionGraphQlService graphQlService() {
+		BatchLoaderRegistry registry = new DefaultBatchLoaderRegistry();
+
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
 		context.register(AuthorController.class);
+		context.registerBean(BatchLoaderRegistry.class, () -> registry);
 		context.refresh();
 
 		AnnotatedControllerConfigurer configurer = new AnnotatedControllerConfigurer();
@@ -145,6 +151,7 @@ public class EntityMappingInvocationTests {
 		return GraphQlSetup.schemaResource(federationSchema)
 				.runtimeWiring(configurer)
 				.schemaFactory(schemaFactory::createGraphQLSchema)
+				.dataLoaders(registry)
 				.toGraphQlService();
 	}
 
@@ -169,10 +176,9 @@ public class EntityMappingInvocationTests {
 			};
 		}
 
-		@SchemaMapping
-		public Author author(Book book) {
-			long id = book.getId();
-			return BookSource.getBook(id).getAuthor();
+		@BatchMapping
+		public Flux<Author> author(List<Book> books) {
+			return Flux.fromIterable(books).map(book -> BookSource.getBook(book.getId()).getAuthor());
 		}
 
 		@GraphQlExceptionHandler