Add WebSocket integration tests w/ Java configuration

Issue: SPR-10835

spring-messaging/src/main/java/org/springframework/messaging/handler/websocket/SubProtocolWebSocketHandler.java

@@ -72,6 +72,7 @@ public class SubProtocolWebSocketHandler implements WebSocketHandler, MessageHan
 		this.outputChannel = outputChannel;
 	}
 
+
 	/**
 	 * Configure one or more handlers to use depending on the sub-protocol requested by
 	 * the client in the WebSocket handshake request.
@@ -130,6 +131,12 @@ public class SubProtocolWebSocketHandler implements WebSocketHandler, MessageHan
 		return this.defaultProtocolHandler;
 	}
 
+	/**
+	 * Return all supported protocols.
+	 */
+	public Set<String> getSupportedProtocols() {
+		return this.protocolHandlers.keySet();
+	}
 
 	@Override
 	public void afterConnectionEstablished(WebSocketSession session) throws Exception {

spring-messaging/src/main/java/org/springframework/messaging/simp/config/StompEndpointRegistration.java

@@ -19,11 +19,13 @@ package org.springframework.messaging.simp.config;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Set;
 
 import org.springframework.messaging.handler.websocket.SubProtocolWebSocketHandler;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
+import org.springframework.util.ObjectUtils;
 import org.springframework.web.HttpRequestHandler;
 import org.springframework.web.socket.server.DefaultHandshakeHandler;
 import org.springframework.web.socket.server.HandshakeHandler;
@@ -46,53 +48,43 @@ public class StompEndpointRegistration {
 
 	private final SubProtocolWebSocketHandler wsHandler;
 
+	private HandshakeHandler handshakeHandler;
+
 	private StompSockJsServiceRegistration sockJsServiceRegistration;
 
-	private TaskScheduler defaultTaskScheduler;
+	private final TaskScheduler defaultSockJsTaskScheduler;
 
 
-	public StompEndpointRegistration(Collection<String> paths, SubProtocolWebSocketHandler webSocketHandler) {
+	public StompEndpointRegistration(Collection<String> paths, SubProtocolWebSocketHandler webSocketHandler,
+			TaskScheduler defaultSockJsTaskScheduler) {
+
 		this.paths = new ArrayList<String>(paths);
 		this.wsHandler = webSocketHandler;
+		this.defaultSockJsTaskScheduler = defaultSockJsTaskScheduler;
 	}
 
 
-	protected List<String> getPaths() {
-		return this.paths;
-	}
-
-	protected SubProtocolWebSocketHandler getSubProtocolWebSocketHandler() {
-		return this.wsHandler;
-	}
-
-	protected StompSockJsServiceRegistration getSockJsServiceRegistration() {
-		return this.sockJsServiceRegistration;
+	public StompEndpointRegistration setHandshakeHandler(HandshakeHandler handshakeHandler) {
+		this.handshakeHandler = handshakeHandler;
+		return this;
 	}
 
 	public SockJsServiceRegistration withSockJS() {
-		this.sockJsServiceRegistration = new StompSockJsServiceRegistration(this.defaultTaskScheduler);
+		this.sockJsServiceRegistration = new StompSockJsServiceRegistration(this.defaultSockJsTaskScheduler);
 		return this.sockJsServiceRegistration;
 	}
 
-	protected void setDefaultTaskScheduler(TaskScheduler defaultTaskScheduler) {
-		this.defaultTaskScheduler = defaultTaskScheduler;
-	}
-
-	protected TaskScheduler getDefaultTaskScheduler() {
-		return this.defaultTaskScheduler;
-	}
-
 	protected MultiValueMap<HttpRequestHandler, String> getMappings() {
 		MultiValueMap<HttpRequestHandler, String> mappings = new LinkedMultiValueMap<HttpRequestHandler, String>();
-		if (getSockJsServiceRegistration() == null) {
-			HandshakeHandler handshakeHandler = createHandshakeHandler();
-			for (String path : getPaths()) {
+		if (this.sockJsServiceRegistration == null) {
+			HandshakeHandler handshakeHandler = getOrCreateHandshakeHandler();
+			for (String path : this.paths) {
 				WebSocketHttpRequestHandler handler = new WebSocketHttpRequestHandler(this.wsHandler, handshakeHandler);
 				mappings.add(handler, path);
 			}
 		}
 		else {
-			SockJsService sockJsService = getSockJsServiceRegistration().getSockJsService();
+			SockJsService sockJsService = this.sockJsServiceRegistration.getSockJsService();
 			for (String path : this.paths) {
 				SockJsHttpRequestHandler httpHandler = new SockJsHttpRequestHandler(sockJsService, this.wsHandler);
 				mappings.add(httpHandler, path.endsWith("/") ? path + "**" : path + "/**");
@@ -101,8 +93,20 @@ public class StompEndpointRegistration {
 		return mappings;
 	}
 
-	protected DefaultHandshakeHandler createHandshakeHandler() {
-		return new DefaultHandshakeHandler();
+	private HandshakeHandler getOrCreateHandshakeHandler() {
+
+		HandshakeHandler handler = (this.handshakeHandler != null)
+				? this.handshakeHandler : new DefaultHandshakeHandler();
+
+		if (handler instanceof DefaultHandshakeHandler) {
+			DefaultHandshakeHandler defaultHandshakeHandler = (DefaultHandshakeHandler) handler;
+			if (ObjectUtils.isEmpty(defaultHandshakeHandler.getSupportedProtocols())) {
+				Set<String> protocols = this.wsHandler.getSupportedProtocols();
+				defaultHandshakeHandler.setSupportedProtocols(protocols.toArray(new String[protocols.size()]));
+			}
+		}
+
+		return handler;
 	}
 
 
@@ -114,7 +118,7 @@ public class StompEndpointRegistration {
 		}
 
 		protected SockJsService getSockJsService() {
-			return super.getSockJsService(getPaths().toArray(new String[getPaths().size()]));
+			return super.getSockJsService(paths.toArray(new String[paths.size()]));
 		}
 	}
 

spring-messaging/src/main/java/org/springframework/messaging/simp/config/StompEndpointRegistry.java

@@ -51,11 +51,11 @@ public class StompEndpointRegistry {
 
 	private int order = 1;
 
-	private TaskScheduler defaultTaskScheduler;
+	private final TaskScheduler defaultSockJsTaskScheduler;
 
 
 	public StompEndpointRegistry(SubProtocolWebSocketHandler webSocketHandler,
-			MutableUserQueueSuffixResolver userQueueSuffixResolver) {
+			MutableUserQueueSuffixResolver userQueueSuffixResolver, TaskScheduler defaultSockJsTaskScheduler) {
 
 		Assert.notNull(webSocketHandler);
 		Assert.notNull(userQueueSuffixResolver);
@@ -63,25 +63,18 @@ public class StompEndpointRegistry {
 		this.wsHandler = webSocketHandler;
 		this.stompHandler = new StompProtocolHandler();
 		this.stompHandler.setUserQueueSuffixResolver(userQueueSuffixResolver);
+		this.defaultSockJsTaskScheduler = defaultSockJsTaskScheduler;
 	}
 
 
 	public StompEndpointRegistration addEndpoint(String... paths) {
 		this.wsHandler.addProtocolHandler(this.stompHandler);
-		StompEndpointRegistration r = new StompEndpointRegistration(Arrays.asList(paths), this.wsHandler);
-		r.setDefaultTaskScheduler(getDefaultTaskScheduler());
+		StompEndpointRegistration r = new StompEndpointRegistration(
+				Arrays.asList(paths), this.wsHandler, this.defaultSockJsTaskScheduler);
 		this.registrations.add(r);
 		return r;
 	}
 
-	protected SubProtocolWebSocketHandler getSubProtocolWebSocketHandler() {
-		return this.wsHandler;
-	}
-
-	protected StompProtocolHandler getStompProtocolHandler() {
-		return this.stompHandler;
-	}
-
 	/**
 	 * Specify the order to use for the STOMP endpoint {@link HandlerMapping} relative to
 	 * other handler mappings configured in the Spring MVC configuration. The default
@@ -91,18 +84,6 @@ public class StompEndpointRegistry {
 		this.order = order;
 	}
 
-	protected int getOrder() {
-		return this.order;
-	}
-
-	protected void setDefaultTaskScheduler(TaskScheduler defaultTaskScheduler) {
-		this.defaultTaskScheduler = defaultTaskScheduler;
-	}
-
-	protected TaskScheduler getDefaultTaskScheduler() {
-		return this.defaultTaskScheduler;
-	}
-
 	/**
 	 * Returns a handler mapping with the mapped ViewControllers; or {@code null} in case of no registrations.
 	 */

spring-messaging/src/main/java/org/springframework/messaging/simp/config/WebSocketMessageBrokerConfigurationSupport.java

@@ -33,6 +33,7 @@ import org.springframework.messaging.support.converter.MessageConverter;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.web.servlet.HandlerMapping;
+import org.springframework.web.socket.server.config.SockJsServiceRegistration;
 
 
 /**
@@ -54,9 +55,8 @@ public abstract class WebSocketMessageBrokerConfigurationSupport {
 
 	@Bean
 	public HandlerMapping brokerWebSocketHandlerMapping() {
-		StompEndpointRegistry registry =
-				new StompEndpointRegistry(subProtocolWebSocketHandler(), userQueueSuffixResolver());
-		registry.setDefaultTaskScheduler(brokerDefaultSockJsTaskScheduler());
+		StompEndpointRegistry registry = new StompEndpointRegistry(
+				subProtocolWebSocketHandler(), userQueueSuffixResolver(), brokerDefaultSockJsTaskScheduler());
 		registerStompEndpoints(registry);
 		return registry.getHandlerMapping();
 	}
@@ -73,11 +73,14 @@ public abstract class WebSocketMessageBrokerConfigurationSupport {
 		return new SimpleUserQueueSuffixResolver();
 	}
 
+	/**
+	 * The default TaskScheduler to use if none is configured via
+	 * {@link SockJsServiceRegistration#setTaskScheduler()}
+	 */
 	@Bean
 	public ThreadPoolTaskScheduler brokerDefaultSockJsTaskScheduler() {
 		ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
 		scheduler.setThreadNamePrefix("BrokerSockJS-");
-		scheduler.setPoolSize(10);
 		return scheduler;
 	}
 
@@ -97,9 +100,7 @@ public abstract class WebSocketMessageBrokerConfigurationSupport {
 	@Bean
 	public ThreadPoolTaskExecutor webSocketChannelExecutor() {
 		ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
-		executor.setCorePoolSize(4);
-		executor.setCorePoolSize(8);
-		executor.setThreadNamePrefix("MessageChannel-");
+		executor.setThreadNamePrefix("BrokerWebSocketChannel-");
 		return executor;
 	}