From fcb11a4ee744385c26fc2a102a67358080259518 Mon Sep 17 00:00:00 2001
From: rstoyanchev <rstoyanchev@vmware.com>
Date: Thu, 14 Sep 2023 20:04:01 +0100
Subject: [PATCH] Fix test failures

Adjust tests with Principal to express nullability or
to provide it. In addition, Mono should not be treated
as a required parameter.

Closes gh-790
---
 .../PrincipalMethodArgumentResolver.java        |  6 +++++-
 ...nticationPrincipalArgumentResolverTests.java |  3 ++-
 .../support/DataFetcherHandlerMethodTests.java  | 17 +++++++++++++----
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/PrincipalMethodArgumentResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/PrincipalMethodArgumentResolver.java
index df7f71f9..eb989c90 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/PrincipalMethodArgumentResolver.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/PrincipalMethodArgumentResolver.java
@@ -64,7 +64,7 @@ public class PrincipalMethodArgumentResolver implements HandlerMethodArgumentRes
 		Mono<Authentication> authMono =
 				ReactiveSecurityContextHolder.getContext().mapNotNull(SecurityContext::getAuthentication);
 
-		if (!parameter.isOptional()) {
+		if (isRequired(parameter)) {
 			authMono = authMono.switchIfEmpty(
 					Mono.error(new AuthenticationCredentialsNotFoundException("No Authentication")));
 		}
@@ -72,4 +72,8 @@ public class PrincipalMethodArgumentResolver implements HandlerMethodArgumentRes
 		return authMono;
 	}
 
+	private static boolean isRequired(MethodParameter parameter) {
+		return (!parameter.isOptional() && !Mono.class.isAssignableFrom(parameter.getParameterType()));
+	}
+
 }
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolverTests.java
index fb0536d1..ce8fd0ba 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolverTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolverTests.java
@@ -33,6 +33,7 @@ import org.springframework.core.ResolvableType;
 import org.springframework.core.annotation.SynthesizingMethodParameter;
 import org.springframework.graphql.data.method.annotation.Argument;
 import org.springframework.graphql.data.method.annotation.QueryMapping;
+import org.springframework.lang.Nullable;
 import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.annotation.AuthenticationPrincipal;
@@ -322,7 +323,7 @@ class AuthenticationPrincipalArgumentResolverTests {
 		}
 
 		@QueryMapping
-		public UserDetails currentUser(@CurrentUser UserDetails userDetails) {
+		public UserDetails currentUser(@CurrentUser @Nullable UserDetails userDetails) {
 			return userDetails;
 		}
 
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java
index 3c43c9e3..d4e06b7d 100644
--- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java
+++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java
@@ -36,7 +36,10 @@ import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComp
 import org.springframework.graphql.data.method.annotation.Argument;
 import org.springframework.graphql.data.method.annotation.QueryMapping;
 import org.springframework.lang.Nullable;
+import org.springframework.security.authentication.TestingAuthenticationToken;
 import org.springframework.security.core.annotation.AuthenticationPrincipal;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextImpl;
 import org.springframework.security.core.userdetails.User;
 import org.springframework.util.ClassUtils;
 
@@ -93,16 +96,22 @@ public class DataFetcherHandlerMethodTests {
 
 		HandlerMethodArgumentResolverComposite resolvers = new HandlerMethodArgumentResolverComposite();
 		resolvers.addResolver(new AuthenticationPrincipalArgumentResolver((beanName, context) -> null));
-		resolvers.addResolver(new ArgumentMethodArgumentResolver(new GraphQlArgumentBinder()));
 
 		DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
 				handlerMethodFor(new TestController(), "handleAndReturnFuture"), resolvers,
 				null, null, false);
 
-		Object result = handlerMethod.invoke(DataFetchingEnvironmentImpl.newDataFetchingEnvironment().build());
+		SecurityContextHolder.setContext(new SecurityContextImpl(new TestingAuthenticationToken("usr", "pwd")));
+		try {
+			Object result = handlerMethod.invoke(
+					DataFetchingEnvironmentImpl.newDataFetchingEnvironment().build());
 
-		assertThat(result).isInstanceOf(Mono.class);
-		assertThat(((Mono<String>) result).block()).isEqualTo("B");
+			assertThat(result).isInstanceOf(Mono.class);
+			assertThat(((Mono<String>) result).block()).isEqualTo("B");
+		}
+		finally {
+			SecurityContextHolder.clearContext();
+		}
 	}
 
 	private static HandlerMethod handlerMethodFor(Object controller, String methodName) {