From 0748a129dc678446855c269468f7c1430dcde43b Mon Sep 17 00:00:00 2001
From: rstoyanchev <rossen.stoyanchev@broadcom.com>
Date: Thu, 2 May 2024 16:10:19 +0100
Subject: [PATCH] Improve union support of SchemaMappingInspector

If the union controller method returns Object, as a fallback
try the package of the controller.

See gh-960
---
 .../execution/SchemaMappingInspector.java     |  8 +++++++
 .../SchemaMappingInspectorUnionTests.java     | 22 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/SchemaMappingInspector.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/SchemaMappingInspector.java
index 2e99455c..2024f983 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/execution/SchemaMappingInspector.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/SchemaMappingInspector.java
@@ -49,6 +49,7 @@ import org.apache.commons.logging.LogFactory;
 
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.BeansException;
+import org.springframework.core.MethodParameter;
 import org.springframework.core.ReactiveAdapter;
 import org.springframework.core.ReactiveAdapterRegistry;
 import org.springframework.core.ResolvableType;
@@ -522,6 +523,13 @@ public final class SchemaMappingInspector {
 							int index = clazz.getName().indexOf(clazz.getSimpleName());
 							resolver.addClassPrefix(outputTypeName, clazz.getName().substring(0, index));
 						}
+						else if (fieldEntry.getValue() instanceof SelfDescribingDataFetcher<?> sddf) {
+							if (sddf.getReturnType().getSource() instanceof MethodParameter param) {
+								clazz = param.getDeclaringClass();
+								int index = clazz.getName().indexOf(clazz.getSimpleName());
+								resolver.addClassPrefix(outputTypeName, clazz.getName().substring(0, index));
+							}
+						}
 					}
 				}
 			}
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/SchemaMappingInspectorUnionTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/SchemaMappingInspectorUnionTests.java
index be337daa..32350a1f 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/execution/SchemaMappingInspectorUnionTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/SchemaMappingInspectorUnionTests.java
@@ -51,7 +51,7 @@ public class SchemaMappingInspectorUnionTests extends SchemaMappingInspectorTest
 	class InterfaceFieldsNotOnJavaInterface {
 
 		@Test
-		void reportUnmappedFields() {
+		void reportUnmappedFieldsByCheckingReturnTypePackage() {
 			SchemaReport report = inspectSchema(schema, SearchController.class);
 			assertThatReport(report)
 					.hasSkippedTypeCount(0)
@@ -60,6 +60,16 @@ public class SchemaMappingInspectorUnionTests extends SchemaMappingInspectorTest
 					.containsUnmappedFields("Video", "title");
 		}
 
+		@Test
+		void reportUnmappedFieldsByCheckingControllerTypePackage() {
+			SchemaReport report = inspectSchema(schema, ObjectSearchController.class);
+			assertThatReport(report)
+					.hasSkippedTypeCount(0)
+					.hasUnmappedFieldCount(3)
+					.containsUnmappedFields("Photo", "height", "width")
+					.containsUnmappedFields("Video", "title");
+		}
+
 
 		sealed interface ResultItem permits Photo, Video { }
 		record Photo() implements ResultItem { }
@@ -73,6 +83,16 @@ public class SchemaMappingInspectorUnionTests extends SchemaMappingInspectorTest
 				throw new UnsupportedOperationException();
 			}
 		}
+
+
+		@Controller
+		static class ObjectSearchController {
+
+			@QueryMapping
+			List<Object> search() {
+				throw new UnsupportedOperationException();
+			}
+		}
 	}