spring-projectsGH-3869: Redesign the test for the advice.

Author: adel-haidar

spring-integration-file/src/main/java/org/springframework/integration/file/remote/session/ContextHolderRequestHandlerAdvice.java

@@ -35,6 +35,19 @@ public class ContextHolderRequestHandlerAdvice extends AbstractRequestHandlerAdv
 
 	public Consumer<Object> contextClearHook;
 
+	/**
+	 * Constructor for the ContextHolderRequestHandlerAdvice.
+	 *
+	 * @param keyProvider The key provider function.
+	 * @param contextSetHook The context set hook function.
+	 * @param contextClearHook The context clear hook function.
+	 */
+	public ContextHolderRequestHandlerAdvice(final Function<Message<?>, Object> keyProvider, final Consumer<Object> contextSetHook, final Consumer<Object> contextClearHook) {
+		this.keyProvider = keyProvider;
+		this.contextSetHook = contextSetHook;
+		this.contextClearHook = contextClearHook;
+	}
+
 	/**
 	 * Returns the key provider function.
 	 * @return The key provider function.
@@ -90,12 +103,12 @@ protected Object doInvoke(ExecutionCallback callback, Object target, Message<?>
 		Object result;
 		final Object key = this.keyProvider.apply(message);
 		try {
-			result = callback.execute();
 			setContext(key);
+			result = callback.execute();
 		}
 		catch (Exception e) {
 			result = e;
-			logger.error("Failure setting context holder for " + message + ": " + e.getMessage());
+			logger.error("Failure setting context for " + message + ": " + e.getMessage());
 		}
 		finally {
 			contextClearHook(key);

spring-integration-file/src/test/java/org/springframework/integration/file/remote/session/ContextHolderRequestHandlerAdviceTests.java

@@ -18,6 +18,8 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Function;
 
 import org.junit.jupiter.api.Test;
 
@@ -53,32 +55,24 @@ public class ContextHolderRequestHandlerAdviceTests {
 	TestSessionFactory bar;
 
 	@Autowired
-	ContextHolderRequestHandlerAdvice advice;
-
-	@Autowired
-	MessageChannel in;
-
-	@Autowired
-	PollableChannel out;
-
-	@Autowired
-	DefaultSessionFactoryLocator<String> sessionFactoryLocator;
-
-	@Autowired
-	TestService testService;
+	Config config;
 
 	@Test
 	public void testFlow() {
-		final Map<String, Object> headers = Map.of("foo", "foo");
-		this.in.send(new GenericMessage<>("foo", headers));
-		Message<?> received = out.receive(0);
+		final Map<String, Object> headers = Map.of("test-key", "foo");
+		config.in().send(new GenericMessage<>("any-payload", headers));
+		Message<?> received = config.out().receive(0);
 		assertThat(received).isNotNull();
+		assertThat(Config.testSessionDuringHandling).isNotNull().isEqualTo(foo.getSession());
+		assertThat(config.dsf().getSession()).isEqualTo(bar.getSession());
 	}
 
 	@Configuration
 	@EnableIntegration
 	public static class Config {
 
+		public static Session<String> testSessionDuringHandling = null;
+
 		@Bean
 		MessageChannel in() {
 			return MessageChannels.direct("in").get();
@@ -107,11 +101,13 @@ TestSessionFactory bar() {
 		@Bean
 		ContextHolderRequestHandlerAdvice advice() {
 			final DelegatingSessionFactory<String> dsf = dsf();
-			final var contextHolderRequestHandlerAdvice = new ContextHolderRequestHandlerAdvice();
-			contextHolderRequestHandlerAdvice.setKeyProvider(m -> m.getHeaders().get("foo"));
-			contextHolderRequestHandlerAdvice.setContextSetHook(dsf::setThreadKey);
-			contextHolderRequestHandlerAdvice.setContextClearHook(c -> dsf.clearThreadKey());
-			return contextHolderRequestHandlerAdvice;
+			final Function<Message<?>, Object> keyProviderFn = m -> m.getHeaders().get("test-key");
+			final Consumer<Object> contextSetHookFn = key -> {
+				dsf.setThreadKey(key);
+				testSessionDuringHandling = dsf.getSession();
+			};
+			final Consumer<Object> contextClearHookFn = key -> dsf.clearThreadKey();
+			return new ContextHolderRequestHandlerAdvice(keyProviderFn, contextSetHookFn, contextClearHookFn);
 		}
 
 		@Bean
@@ -134,9 +130,6 @@ public SessionFactoryLocator<String> sessionFactoryLocator() {
 	@Component
 	public static class TestService {
 
-		@Autowired
-		private ContextHolderRequestHandlerAdvice advice;
-
 		@ServiceActivator(inputChannel = "in", outputChannel = "out", adviceChain = "advice")
 		public String test(Message<String> message) {
 			return message.getPayload();