WebSocketMessageBrokerConfigurer allows to configure Lifecycle phase

Closes gh-32205

spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.function.Consumer;
 
 import org.junit.jupiter.api.Test;
 
@@ -35,6 +36,7 @@ import org.springframework.messaging.handler.annotation.MessageMapping;
 import org.springframework.messaging.handler.annotation.SendTo;
 import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.simp.annotation.SubscribeMapping;
+import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler;
 import org.springframework.messaging.simp.broker.DefaultSubscriptionRegistry;
 import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
@@ -45,6 +47,7 @@ import org.springframework.messaging.support.AbstractSubscribableChannel;
 import org.springframework.messaging.support.ChannelInterceptor;
 import org.springframework.messaging.support.ExecutorSubscribableChannel;
 import org.springframework.messaging.support.ImmutableMessageChannelInterceptor;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.stereotype.Controller;
 import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
@@ -57,6 +60,7 @@ import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.messaging.StompTextMessageBuilder;
 import org.springframework.web.socket.messaging.SubProtocolHandler;
 import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
+import org.springframework.web.socket.server.support.WebSocketHandlerMapping;
 import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -207,6 +211,27 @@ class WebSocketMessageBrokerConfigurationSupportTests {
 		assertThat(actual).matches(expected);
 	}
 
+	@Test
+	void lifecyclePhase() {
+		ApplicationContext context = createContext(LifecyclePhaseConfig.class);
+
+		int phase = 99;
+		Consumer<String> executorTester = beanName ->
+				assertThat(context.getBean(beanName, ThreadPoolTaskExecutor.class).getPhase()).isEqualTo(phase);
+
+		executorTester.accept("clientInboundChannelExecutor");
+		executorTester.accept("clientOutboundChannelExecutor");
+		executorTester.accept("brokerChannelExecutor");
+
+		assertThat(context.getBean(SimpAnnotationMethodMessageHandler.class).getPhase()).isEqualTo(phase);
+		assertThat(context.getBean(UserDestinationMessageHandler.class).getPhase()).isEqualTo(phase);
+		assertThat(context.getBean(ThreadPoolTaskScheduler.class).getPhase()).isEqualTo(phase);
+		assertThat(context.getBean(WebSocketHandlerMapping.class).getPhase()).isEqualTo(phase);
+
+		assertThat(context.getBean(SubProtocolWebSocketHandler.class).getPhase()).isEqualTo(phase);
+		assertThat(context.getBean(SimpleBrokerMessageHandler.class).getPhase()).isEqualTo(phase);
+	}
+
 	@Test
 	void webSocketHandlerDecorator() throws Exception {
 		ApplicationContext context = createContext(WebSocketHandlerDecoratorConfig.class);
@@ -347,4 +372,30 @@ class WebSocketMessageBrokerConfigurationSupportTests {
 		}
 	}
 
+
+	@Configuration
+	static class LifecyclePhaseConfig extends DelegatingWebSocketMessageBrokerConfiguration {
+
+		@Bean
+		public WebSocketMessageBrokerConfigurer getConfigurer() {
+			return new WebSocketMessageBrokerConfigurer() {
+
+				@Override
+				public void registerStompEndpoints(StompEndpointRegistry registry) {
+					registry.addEndpoint("/broker");
+				}
+
+				@Override
+				public void configureMessageBroker(MessageBrokerRegistry registry) {
+					registry.enableSimpleBroker();
+				}
+
+				@Override
+				public Integer getPhase() {
+					return 99;
+				}
+			};
+		}
+	}
+
 }