From 1648151dd232973bacda2bf624f260251c553982 Mon Sep 17 00:00:00 2001
From: Marcus Da Coregio <marcusdacoregio@gmail.com>
Date: Thu, 17 Nov 2022 10:08:40 -0300
Subject: [PATCH] Register hints for @WithSecurityContext on class level

Issue gh-12215
---
 .../hint/WithSecurityContextTestRuntimeHints.java   | 13 +++++++++++--
 .../WithSecurityContextTestRuntimeHintsTests.java   |  9 +++++++++
 .../test/context/showcase/WithMockCustomUser.java   |  4 ++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/test/src/main/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHints.java b/test/src/main/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHints.java
index 5b41094dce..0f50bd8368 100644
--- a/test/src/main/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHints.java
+++ b/test/src/main/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHints.java
@@ -17,6 +17,7 @@
 package org.springframework.security.test.aot.hint;
 
 import java.util.Arrays;
+import java.util.stream.Stream;
 
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
@@ -38,13 +39,21 @@ class WithSecurityContextTestRuntimeHints implements TestRuntimeHintsRegistrar {
 
 	@Override
 	public void registerHints(RuntimeHints hints, Class<?> testClass, ClassLoader classLoader) {
-		Arrays.stream(testClass.getDeclaredMethods())
-				.map((method) -> MergedAnnotations.from(method, SUPERCLASS).get(WithSecurityContext.class))
+		Stream.concat(getClassAnnotations(testClass), getMethodAnnotations(testClass))
 				.filter(MergedAnnotation::isPresent)
 				.map((withSecurityContext) -> withSecurityContext.getClass("factory"))
 				.forEach((factory) -> registerDeclaredConstructors(hints, factory));
 	}
 
+	private Stream<MergedAnnotation<WithSecurityContext>> getClassAnnotations(Class<?> testClass) {
+		return MergedAnnotations.search(SUPERCLASS).from(testClass).stream(WithSecurityContext.class);
+	}
+
+	private Stream<MergedAnnotation<WithSecurityContext>> getMethodAnnotations(Class<?> testClass) {
+		return Arrays.stream(testClass.getDeclaredMethods())
+				.map((method) -> MergedAnnotations.from(method, SUPERCLASS).get(WithSecurityContext.class));
+	}
+
 	private void registerDeclaredConstructors(RuntimeHints hints, Class<?> factory) {
 		hints.reflection().registerType(factory, MemberCategory.INVOKE_DECLARED_CONSTRUCTORS);
 	}
diff --git a/test/src/test/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHintsTests.java b/test/src/test/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHintsTests.java
index 1eb5fc3f3f..f8db29d8bb 100644
--- a/test/src/test/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHintsTests.java
+++ b/test/src/test/java/org/springframework/security/test/aot/hint/WithSecurityContextTestRuntimeHintsTests.java
@@ -28,6 +28,8 @@ import org.springframework.aot.hint.TypeReference;
 import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
 import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.test.context.showcase.WithMockCustomUser;
+import org.springframework.security.test.context.showcase.WithMockCustomUserSecurityContextFactory;
 import org.springframework.security.test.context.support.WithAnonymousUser;
 import org.springframework.security.test.context.support.WithMockUser;
 import org.springframework.security.test.context.support.WithSecurityContext;
@@ -39,6 +41,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 /**
  * Tests for {@link WithSecurityContextTestRuntimeHints}.
  */
+@WithMockCustomUser
 class WithSecurityContextTestRuntimeHintsTests {
 
 	private final RuntimeHints hints = new RuntimeHints();
@@ -85,6 +88,12 @@ class WithSecurityContextTestRuntimeHintsTests {
 				.withMemberCategory(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS)).accepts(this.hints);
 	}
 
+	@Test
+	void withMockCustomUserOnClassHasHints() {
+		assertThat(RuntimeHintsPredicates.reflection().onType(WithMockCustomUserSecurityContextFactory.class)
+				.withMemberCategory(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS)).accepts(this.hints);
+	}
+
 	@Retention(RetentionPolicy.RUNTIME)
 	@WithSecurityContext(factory = WithMockTestUserSecurityContextFactory.class)
 	@interface WithMockTestUser {
diff --git a/test/src/test/java/org/springframework/security/test/context/showcase/WithMockCustomUser.java b/test/src/test/java/org/springframework/security/test/context/showcase/WithMockCustomUser.java
index a3c61627f2..fcfe636315 100644
--- a/test/src/test/java/org/springframework/security/test/context/showcase/WithMockCustomUser.java
+++ b/test/src/test/java/org/springframework/security/test/context/showcase/WithMockCustomUser.java
@@ -16,11 +16,15 @@
 
 package org.springframework.security.test.context.showcase;
 
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
 import org.springframework.security.test.context.support.WithSecurityContext;
 
 /**
  * @author Rob Winch
  */
+@Retention(RetentionPolicy.RUNTIME)
 @WithSecurityContext(factory = WithMockCustomUserSecurityContextFactory.class)
 public @interface WithMockCustomUser {