Polish WebSocket/STOMP Java config

Ensure configuration provided for WebSocketHandler's (eg interceptors,
or HandshakeHandler) are passed on to the SockJsService if congiured.

Better separate Servlet-specific parts of the configuration to make it
more obvious where non-Servlet alternatives could fit in.

Add more tests.

Improve WebSocket integration tests.

spring-websocket/src/main/java/org/springframework/web/socket/server/config/AbstractWebSocketHandlerRegistration.java

@@ -0,0 +1,153 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.socket.server.config;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.util.Assert;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.server.DefaultHandshakeHandler;
+import org.springframework.web.socket.server.HandshakeHandler;
+import org.springframework.web.socket.server.HandshakeInterceptor;
+import org.springframework.web.socket.sockjs.SockJsService;
+import org.springframework.web.socket.sockjs.transport.handler.WebSocketTransportHandler;
+
+
+/**
+ * Base class for {@link WebSocketHandlerRegistration}s that gathers all the configuration
+ * options but allows sub-classes to put together the actual HTTP request mappings.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public abstract class AbstractWebSocketHandlerRegistration<M> implements WebSocketHandlerRegistration {
+
+	private MultiValueMap<WebSocketHandler, String> handlerMap = new LinkedMultiValueMap<WebSocketHandler, String>();
+
+	private HandshakeInterceptor[] interceptors;
+
+	private HandshakeHandler handshakeHandler;
+
+	private SockJsServiceRegistration sockJsServiceRegistration;
+
+	private final TaskScheduler sockJsTaskScheduler;
+
+
+	public AbstractWebSocketHandlerRegistration(TaskScheduler defaultTaskScheduler) {
+		this.sockJsTaskScheduler = defaultTaskScheduler;
+	}
+
+	@Override
+	public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String... paths) {
+		Assert.notNull(handler);
+		Assert.notEmpty(paths);
+		this.handlerMap.put(handler, Arrays.asList(paths));
+		return this;
+	}
+
+	@Override
+	public WebSocketHandlerRegistration setHandshakeHandler(HandshakeHandler handshakeHandler) {
+		this.handshakeHandler = handshakeHandler;
+		return this;
+	}
+
+	public HandshakeHandler getHandshakeHandler() {
+		return handshakeHandler;
+	}
+
+	@Override
+	public WebSocketHandlerRegistration addInterceptors(HandshakeInterceptor... interceptors) {
+		this.interceptors = interceptors;
+		return this;
+	}
+
+	protected HandshakeInterceptor[] getInterceptors() {
+		return this.interceptors;
+	}
+
+	/**
+	 * @param interceptors the interceptors to set
+	 */
+	public void setInterceptors(HandshakeInterceptor[] interceptors) {
+		this.interceptors = interceptors;
+	}
+
+	@Override
+	public SockJsServiceRegistration withSockJS() {
+
+		this.sockJsServiceRegistration = new SockJsServiceRegistration(this.sockJsTaskScheduler);
+		this.sockJsServiceRegistration.setInterceptors(this.interceptors);
+
+		if (this.handshakeHandler != null) {
+			WebSocketTransportHandler transportHandler = new WebSocketTransportHandler(this.handshakeHandler);
+			this.sockJsServiceRegistration.setTransportHandlerOverrides(transportHandler);
+		}
+
+		return this.sockJsServiceRegistration;
+	}
+
+	final M getMappings() {
+
+		M mappings = createMappings();
+
+		if (this.sockJsServiceRegistration != null) {
+			SockJsService sockJsService = this.sockJsServiceRegistration.getSockJsService(getAllPrefixes());
+			for (WebSocketHandler wsHandler : this.handlerMap.keySet()) {
+				for (String path : this.handlerMap.get(wsHandler)) {
+					String pathPattern = path.endsWith("/") ? path + "**" : path + "/**";
+					addSockJsServiceMapping(mappings, sockJsService, wsHandler, pathPattern);
+				}
+			}
+		}
+		else {
+			HandshakeHandler handshakeHandler = getOrCreateHandshakeHandler();
+			for (WebSocketHandler wsHandler : this.handlerMap.keySet()) {
+				for (String path : this.handlerMap.get(wsHandler)) {
+					addWebSocketHandlerMapping(mappings, wsHandler, handshakeHandler, this.interceptors, path);
+				}
+			}
+		}
+
+		return mappings;
+	}
+
+	private final String[] getAllPrefixes() {
+		List<String> all = new ArrayList<String>();
+		for (List<String> prefixes: this.handlerMap.values()) {
+			all.addAll(prefixes);
+		}
+		return all.toArray(new String[all.size()]);
+	}
+
+	private HandshakeHandler getOrCreateHandshakeHandler() {
+		return (this.handshakeHandler != null) ? this.handshakeHandler : new DefaultHandshakeHandler();
+	}
+
+	protected abstract M createMappings();
+
+	protected abstract void addSockJsServiceMapping(M mappings, SockJsService sockJsService,
+			WebSocketHandler handler, String pathPattern);
+
+	protected abstract void addWebSocketHandlerMapping(M mappings, WebSocketHandler wsHandler,
+			HandshakeHandler handshakeHandler, HandshakeInterceptor[] interceptors, String path);
+
+}

spring-websocket/src/main/java/org/springframework/web/socket/server/config/ServletWebSocketHandlerRegistration.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.socket.server.config;
+
+import java.util.Arrays;
+
+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.WebSocketHandler;
+import org.springframework.web.socket.server.HandshakeHandler;
+import org.springframework.web.socket.server.HandshakeInterceptor;
+import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
+import org.springframework.web.socket.sockjs.SockJsHttpRequestHandler;
+import org.springframework.web.socket.sockjs.SockJsService;
+
+
+/**
+ * A helper class for configuring {@link WebSocketHandler} request handling
+ * including SockJS fallback options.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class ServletWebSocketHandlerRegistration
+		extends AbstractWebSocketHandlerRegistration<MultiValueMap<HttpRequestHandler, String>> {
+
+
+	public ServletWebSocketHandlerRegistration(TaskScheduler sockJsTaskScheduler) {
+		super(sockJsTaskScheduler);
+	}
+
+
+	@Override
+	protected MultiValueMap<HttpRequestHandler, String> createMappings() {
+		return new LinkedMultiValueMap<HttpRequestHandler, String>();
+	}
+
+	@Override
+	protected void addSockJsServiceMapping(MultiValueMap<HttpRequestHandler, String> mappings,
+			SockJsService sockJsService, WebSocketHandler handler, String pathPattern) {
+
+		SockJsHttpRequestHandler httpHandler = new SockJsHttpRequestHandler(sockJsService, handler);
+		mappings.add(httpHandler, pathPattern);
+	}
+
+	@Override
+	protected void addWebSocketHandlerMapping(MultiValueMap<HttpRequestHandler, String> mappings,
+			WebSocketHandler wsHandler, HandshakeHandler handshakeHandler,
+			HandshakeInterceptor[] interceptors, String path) {
+
+		WebSocketHttpRequestHandler httpHandler = new WebSocketHttpRequestHandler(wsHandler, handshakeHandler);
+		if (!ObjectUtils.isEmpty(interceptors)) {
+			httpHandler.setHandshakeInterceptors(Arrays.asList(interceptors));
+		}
+		mappings.add(httpHandler, path);
+	}
+
+}

spring-websocket/src/main/java/org/springframework/web/socket/server/config/ServletWebSocketHandlerRegistry.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.socket.server.config;
+
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.HttpRequestHandler;
+import org.springframework.web.servlet.HandlerMapping;
+import org.springframework.web.servlet.handler.AbstractHandlerMapping;
+import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
+import org.springframework.web.socket.WebSocketHandler;
+
+
+/**
+ * A {@link WebSocketHandlerRegistry} that maps {@link WebSocketHandler}s to URLs for use
+ * in a Servlet container.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public class ServletWebSocketHandlerRegistry implements WebSocketHandlerRegistry {
+
+	private final List<ServletWebSocketHandlerRegistration> registrations =
+			new ArrayList<ServletWebSocketHandlerRegistration>();
+
+	private TaskScheduler sockJsTaskScheduler;
+
+
+	public ServletWebSocketHandlerRegistry(ThreadPoolTaskScheduler sockJsTaskScheduler) {
+		this.sockJsTaskScheduler = sockJsTaskScheduler;
+	}
+
+	@Override
+	public WebSocketHandlerRegistration addHandler(WebSocketHandler webSocketHandler, String... paths) {
+		ServletWebSocketHandlerRegistration r = new ServletWebSocketHandlerRegistration(this.sockJsTaskScheduler);
+		r.addHandler(webSocketHandler, paths);
+		this.registrations.add(r);
+		return r;
+	}
+
+	/**
+	 * Returns a {@link HandlerMapping} with mapped {@link HttpRequestHandler}s.
+	 */
+	AbstractHandlerMapping getHandlerMapping() {
+		Map<String, Object> urlMap = new LinkedHashMap<String, Object>();
+		for (ServletWebSocketHandlerRegistration registration : this.registrations) {
+			MultiValueMap<HttpRequestHandler, String> mappings = registration.getMappings();
+			for (HttpRequestHandler httpHandler : mappings.keySet()) {
+				for (String pattern : mappings.get(httpHandler)) {
+					urlMap.put(pattern, httpHandler);
+				}
+			}
+		}
+		SimpleUrlHandlerMapping hm = new SimpleUrlHandlerMapping();
+		hm.setUrlMap(urlMap);
+		return hm;
+	}
+
+}

spring-websocket/src/main/java/org/springframework/web/socket/server/config/WebSocketConfigurationSupport.java

@@ -19,6 +19,7 @@ package org.springframework.web.socket.server.config;
 import org.springframework.context.annotation.Bean;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.web.servlet.HandlerMapping;
+import org.springframework.web.servlet.handler.AbstractHandlerMapping;
 
 
 /**
@@ -32,9 +33,11 @@ public class WebSocketConfigurationSupport {
 
 	@Bean
 	public HandlerMapping webSocketHandlerMapping() {
-		WebSocketHandlerRegistry registry = new WebSocketHandlerRegistry(defaultSockJsTaskScheduler());
+		ServletWebSocketHandlerRegistry registry = new ServletWebSocketHandlerRegistry(defaultSockJsTaskScheduler());
 		registerWebSocketHandlers(registry);
-		return registry.getHandlerMapping();
+		AbstractHandlerMapping hm = registry.getHandlerMapping();
+		hm.setOrder(1);
+		return hm;
 	}
 
 	protected void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
@@ -42,7 +45,20 @@ public class WebSocketConfigurationSupport {
 
 	/**
 	 * The default TaskScheduler to use if none is configured via
-	 * {@link SockJsServiceRegistration#setTaskScheduler()}
+	 * {@link SockJsServiceRegistration#setTaskScheduler()}, i.e.
+	 * <pre class="code">
+	 * &#064;Configuration
+	 * &#064;EnableWebSocket
+	 * public class WebSocketConfig implements WebSocketConfigurer {
+	 *
+	 *   public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
+	 *     registry.addHandler(myWsHandler(), "/echo").withSockJS().setTaskScheduler(myScheduler());
+	 *   }
+	 *
+	 *   // ...
+	 *
+	 * }
+	 * </pre>
 	 */
 	@Bean
 	public ThreadPoolTaskScheduler defaultSockJsTaskScheduler() {

spring-websocket/src/main/java/org/springframework/web/socket/server/config/WebSocketHandlerRegistration.java

@@ -16,110 +16,37 @@
 
 package org.springframework.web.socket.server.config;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-
-import org.springframework.scheduling.TaskScheduler;
-import org.springframework.util.Assert;
-import org.springframework.util.LinkedMultiValueMap;
-import org.springframework.util.MultiValueMap;
-import org.springframework.web.HttpRequestHandler;
 import org.springframework.web.socket.WebSocketHandler;
-import org.springframework.web.socket.server.DefaultHandshakeHandler;
 import org.springframework.web.socket.server.HandshakeHandler;
 import org.springframework.web.socket.server.HandshakeInterceptor;
-import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
-import org.springframework.web.socket.sockjs.SockJsHttpRequestHandler;
-import org.springframework.web.socket.sockjs.SockJsService;
-
 
 /**
- * A helper class for configuring {@link WebSocketHandler} request handling
- * including SockJS fallback options.
+ * Provides methods for configuring a WebSocket handler.
  *
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public class WebSocketHandlerRegistration {
+public interface WebSocketHandlerRegistration {
 
-	private MultiValueMap<WebSocketHandler, String> handlerMap =
-			new LinkedMultiValueMap<WebSocketHandler, String>();
+	/**
+	 * Add more handlers that will share the same configuration (interceptors, SockJS
+	 * config, etc)
+	 */
+	WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String... paths);
 
-	private final List<HandshakeInterceptor> interceptors = new ArrayList<HandshakeInterceptor>();
+	/**
+	 * Configure interceptors for the handshake request.
+	 */
+	WebSocketHandlerRegistration addInterceptors(HandshakeInterceptor... interceptors);
 
-	private HandshakeHandler handshakeHandler;
+	/**
+	 * Configure the HandshakeHandler to use.
+	 */
+	WebSocketHandlerRegistration setHandshakeHandler(HandshakeHandler handshakeHandler);
 
-	private SockJsServiceRegistration sockJsServiceRegistration;
+	/**
+	 * Enable SockJS fallback options.
+	 */
+	SockJsServiceRegistration withSockJS();
 
-	private final TaskScheduler defaultTaskScheduler;
-
-
-	public WebSocketHandlerRegistration(TaskScheduler defaultTaskScheduler) {
-		this.defaultTaskScheduler = defaultTaskScheduler;
-	}
-
-	public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String... paths) {
-		Assert.notNull(handler);
-		Assert.notEmpty(paths);
-		this.handlerMap.put(handler, Arrays.asList(paths));
-		return this;
-	}
-
-	public WebSocketHandlerRegistration setHandshakeHandler(HandshakeHandler handshakeHandler) {
-		this.handshakeHandler = handshakeHandler;
-		return this;
-	}
-
-	public HandshakeHandler getHandshakeHandler() {
-		return handshakeHandler;
-	}
-
-	public void addInterceptors(HandshakeInterceptor... interceptors) {
-		this.interceptors.addAll(Arrays.asList(interceptors));
-	}
-
-	public SockJsServiceRegistration withSockJS() {
-		this.sockJsServiceRegistration = new SockJsServiceRegistration(this.defaultTaskScheduler);
-		this.sockJsServiceRegistration.setInterceptors(
-				this.interceptors.toArray(new HandshakeInterceptor[this.interceptors.size()]));
-		return this.sockJsServiceRegistration;
-	}
-
-	MultiValueMap<HttpRequestHandler, String> getMappings() {
-		MultiValueMap<HttpRequestHandler, String> mappings = new LinkedMultiValueMap<HttpRequestHandler, String>();
-		if (this.sockJsServiceRegistration == null) {
-			HandshakeHandler handshakeHandler = getOrCreateHandshakeHandler();
-			for (WebSocketHandler handler : this.handlerMap.keySet()) {
-				for (String path : this.handlerMap.get(handler)) {
-					WebSocketHttpRequestHandler httpHandler = new WebSocketHttpRequestHandler(handler, handshakeHandler);
-					httpHandler.setHandshakeInterceptors(this.interceptors);
-					mappings.add(httpHandler, path);
-				}
-			}
-		}
-		else {
-			SockJsService sockJsService = this.sockJsServiceRegistration.getSockJsService(getAllPrefixes());
-			for (WebSocketHandler handler : this.handlerMap.keySet()) {
-				for (String path : this.handlerMap.get(handler)) {
-					SockJsHttpRequestHandler httpHandler = new SockJsHttpRequestHandler(sockJsService, handler);
-					mappings.add(httpHandler, path.endsWith("/") ? path + "**" : path + "/**");
-				}
-			}
-		}
-		return mappings;
-	}
-
-	private HandshakeHandler getOrCreateHandshakeHandler() {
-		return (this.handshakeHandler != null) ? this.handshakeHandler : new DefaultHandshakeHandler();
-	}
-
-	private final String[] getAllPrefixes() {
-		List<String> all = new ArrayList<String>();
-		for (List<String> prefixes: this.handlerMap.values()) {
-			all.addAll(prefixes);
-		}
-		return all.toArray(new String[all.size()]);
-	}
-
-}
+}

spring-websocket/src/main/java/org/springframework/web/socket/server/config/WebSocketHandlerRegistry.java

@@ -16,76 +16,19 @@
 
 package org.springframework.web.socket.server.config;
 
-import java.util.ArrayList;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-
-import org.springframework.scheduling.TaskScheduler;
-import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
-import org.springframework.util.MultiValueMap;
-import org.springframework.web.HttpRequestHandler;
-import org.springframework.web.servlet.HandlerMapping;
-import org.springframework.web.servlet.handler.AbstractHandlerMapping;
-import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.socket.WebSocketHandler;
 
-
 /**
- * A helper class for configuring {@link WebSocketHandler} request handling.
+ * Provides methods for configuring {@link WebSocketHandler} request mappings.
  *
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public class WebSocketHandlerRegistry {
-
-	private final List<WebSocketHandlerRegistration> registrations = new ArrayList<WebSocketHandlerRegistration>();
-
-	private int order = 1;
-
-	private TaskScheduler defaultTaskScheduler;
-
-
-	public WebSocketHandlerRegistry(ThreadPoolTaskScheduler defaultSockJsTaskScheduler) {
-		this.defaultTaskScheduler = defaultSockJsTaskScheduler;
-	}
-
-	public WebSocketHandlerRegistration addHandler(WebSocketHandler wsHandler, String... paths) {
-		WebSocketHandlerRegistration r = new WebSocketHandlerRegistration(this.defaultTaskScheduler);
-		r.addHandler(wsHandler, paths);
-		this.registrations.add(r);
-		return r;
-	}
-
-	protected List<WebSocketHandlerRegistration> getRegistrations() {
-		return this.registrations;
-	}
+public interface WebSocketHandlerRegistry {
 
 	/**
-	 * Specify the order to use for WebSocket {@link HandlerMapping} relative to other
-	 * handler mappings configured in the Spring MVC configuration. The default value is 1.
+	 * Configure a WebSocketHandler at the specified URL paths.
 	 */
-	public void setOrder(int order) {
-		this.order = order;
-	}
+	WebSocketHandlerRegistration addHandler(WebSocketHandler webSocketHandler, String... paths);
 
-	/**
-	 * Returns a handler mapping with the mapped ViewControllers; or {@code null} in case of no registrations.
-	 */
-	AbstractHandlerMapping getHandlerMapping() {
-		Map<String, Object> urlMap = new LinkedHashMap<String, Object>();
-		for (WebSocketHandlerRegistration registration : this.registrations) {
-			MultiValueMap<HttpRequestHandler, String> mappings = registration.getMappings();
-			for (HttpRequestHandler httpHandler : mappings.keySet()) {
-				for (String pattern : mappings.get(httpHandler)) {
-					urlMap.put(pattern, httpHandler);
-				}
-			}
-		}
-		SimpleUrlHandlerMapping hm = new SimpleUrlHandlerMapping();
-		hm.setOrder(this.order);
-		hm.setUrlMap(urlMap);
-		return hm;
-	}
-
-}
+}

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/handler/WebSocketTransportHandler.java

@@ -61,6 +61,10 @@ public class WebSocketTransportHandler extends TransportHandlerSupport
 		return TransportType.WEBSOCKET;
 	}
 
+	public HandshakeHandler getHandshakeHandler() {
+		return this.handshakeHandler;
+	}
+
 	@Override
 	public AbstractSockJsSession createSession(String sessionId, WebSocketHandler wsHandler,
 			Map<String, Object> attributes) {

spring-websocket/src/test/java/org/springframework/web/socket/AbstractWebSocketIntegrationTests.java

@@ -25,6 +25,7 @@ import org.junit.runners.Parameterized.Parameter;
 import org.springframework.context.Lifecycle;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
 import org.springframework.web.socket.client.WebSocketClient;
 import org.springframework.web.socket.server.DefaultHandshakeHandler;
 import org.springframework.web.socket.server.HandshakeHandler;
@@ -52,14 +53,26 @@ public abstract class AbstractWebSocketIntegrationTests {
 	@Parameter(1)
 	public WebSocketClient webSocketClient;
 
+	protected AnnotationConfigWebApplicationContext wac;
+
 
 	@Before
 	public void setup() throws Exception {
+
+		this.wac = new AnnotationConfigWebApplicationContext();
+		this.wac.register(getAnnotatedConfigClasses());
+		this.wac.register(upgradeStrategyConfigTypes.get(this.server.getClass()));
+
 		if (this.webSocketClient instanceof Lifecycle) {
 			((Lifecycle) this.webSocketClient).start();
 		}
+
+		this.server.init(this.wac);
+		this.server.start();
 	}
 
+	protected abstract Class<?>[] getAnnotatedConfigClasses();
+
 	@After
 	public void teardown() throws Exception {
 		try {
@@ -76,10 +89,6 @@ public abstract class AbstractWebSocketIntegrationTests {
 		return "ws://localhost:" + this.server.getPort();
 	}
 
-	protected Class<?> getUpgradeStrategyConfigClass() {
-		return upgradeStrategyConfigTypes.get(this.server.getClass());
-	}
-
 
 	static abstract class AbstractRequestUpgradeStrategyConfig {
 

spring-websocket/src/test/java/org/springframework/web/socket/server/config/AbstractWebSocketHandlerRegistrationTests.java

@@ -0,0 +1,210 @@
+/*
+ * Copyright 2002-2013 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.socket.server.config;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.adapter.TextWebSocketHandlerAdapter;
+import org.springframework.web.socket.server.DefaultHandshakeHandler;
+import org.springframework.web.socket.server.HandshakeHandler;
+import org.springframework.web.socket.server.HandshakeInterceptor;
+import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;
+import org.springframework.web.socket.sockjs.SockJsService;
+import org.springframework.web.socket.sockjs.transport.TransportType;
+import org.springframework.web.socket.sockjs.transport.handler.DefaultSockJsService;
+import org.springframework.web.socket.sockjs.transport.handler.WebSocketTransportHandler;
+
+import static org.junit.Assert.*;
+
+
+/**
+ * Test fixture for {@link AbstractWebSocketHandlerRegistration}.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class AbstractWebSocketHandlerRegistrationTests {
+
+	private TestWebSocketHandlerRegistration registration;
+
+	private TaskScheduler taskScheduler;
+
+
+	@Before
+	public void setup() {
+		this.taskScheduler = Mockito.mock(TaskScheduler.class);
+		this.registration = new TestWebSocketHandlerRegistration(taskScheduler);
+	}
+
+	@Test
+	public void minimal() {
+
+		WebSocketHandler wsHandler = new TextWebSocketHandlerAdapter();
+		this.registration.addHandler(wsHandler, "/foo", "/bar");
+
+		List<Mapping> mappings = this.registration.getMappings();
+		assertEquals(2, mappings.size());
+
+		Mapping m1 = mappings.get(0);
+		assertEquals(wsHandler, m1.webSocketHandler);
+		assertEquals("/foo", m1.path);
+
+		Mapping m2 = mappings.get(1);
+		assertEquals(wsHandler, m2.webSocketHandler);
+		assertEquals("/bar", m2.path);
+	}
+
+	@Test
+	public void interceptors() {
+
+		WebSocketHandler wsHandler = new TextWebSocketHandlerAdapter();
+		HttpSessionHandshakeInterceptor interceptor = new HttpSessionHandshakeInterceptor();
+
+		this.registration.addHandler(wsHandler, "/foo").addInterceptors(interceptor);
+
+		List<Mapping> mappings = this.registration.getMappings();
+		assertEquals(1, mappings.size());
+
+		Mapping m1 = mappings.get(0);
+		assertEquals(wsHandler, m1.webSocketHandler);
+		assertEquals("/foo", m1.path);
+		assertArrayEquals(new HandshakeInterceptor[] { interceptor }, m1.interceptors);
+	}
+
+	@Test
+	public void interceptorsPassedToSockJsRegistration() {
+
+		WebSocketHandler wsHandler = new TextWebSocketHandlerAdapter();
+		HttpSessionHandshakeInterceptor interceptor = new HttpSessionHandshakeInterceptor();
+
+		this.registration.addHandler(wsHandler, "/foo").addInterceptors(interceptor).withSockJS();
+
+		List<Mapping> mappings = this.registration.getMappings();
+		assertEquals(1, mappings.size());
+
+		Mapping m1 = mappings.get(0);
+		assertEquals(wsHandler, m1.webSocketHandler);
+		assertEquals("/foo/**", m1.path);
+		assertNotNull(m1.sockJsService);
+		assertEquals(Arrays.asList(interceptor), m1.sockJsService.getHandshakeInterceptors());
+	}
+
+	@Test
+	public void handshakeHandler() {
+
+		WebSocketHandler wsHandler = new TextWebSocketHandlerAdapter();
+		HandshakeHandler handshakeHandler = new DefaultHandshakeHandler();
+
+		this.registration.addHandler(wsHandler, "/foo").setHandshakeHandler(handshakeHandler);
+
+		List<Mapping> mappings = this.registration.getMappings();
+		assertEquals(1, mappings.size());
+
+		Mapping m1 = mappings.get(0);
+		assertEquals(wsHandler, m1.webSocketHandler);
+		assertEquals("/foo", m1.path);
+		assertSame(handshakeHandler, m1.handshakeHandler);
+	}
+
+	@Test
+	public void handshakeHandlerPassedToSockJsRegistration() {
+
+		WebSocketHandler wsHandler = new TextWebSocketHandlerAdapter();
+		HandshakeHandler handshakeHandler = new DefaultHandshakeHandler();
+
+		this.registration.addHandler(wsHandler, "/foo").setHandshakeHandler(handshakeHandler).withSockJS();
+
+		List<Mapping> mappings = this.registration.getMappings();
+		assertEquals(1, mappings.size());
+
+		Mapping m1 = mappings.get(0);
+		assertEquals(wsHandler, m1.webSocketHandler);
+		assertEquals("/foo/**", m1.path);
+		assertNotNull(m1.sockJsService);
+
+		WebSocketTransportHandler transportHandler =
+				(WebSocketTransportHandler) m1.sockJsService.getTransportHandlers().get(TransportType.WEBSOCKET);
+		assertSame(handshakeHandler, transportHandler.getHandshakeHandler());
+	}
+
+
+	private static class TestWebSocketHandlerRegistration  extends AbstractWebSocketHandlerRegistration<List<Mapping>> {
+
+
+		public TestWebSocketHandlerRegistration(TaskScheduler sockJsTaskScheduler) {
+			super(sockJsTaskScheduler);
+		}
+
+		@Override
+		protected List<Mapping> createMappings() {
+			return new ArrayList<>();
+		}
+
+		@Override
+		protected void addSockJsServiceMapping(List<Mapping> mappings, SockJsService sockJsService,
+				WebSocketHandler wsHandler, String pathPattern) {
+
+			mappings.add(new Mapping(wsHandler, pathPattern, sockJsService));
+		}
+
+		@Override
+		protected void addWebSocketHandlerMapping(List<Mapping> mappings,
+				WebSocketHandler wsHandler, HandshakeHandler handshakeHandler,
+				HandshakeInterceptor[] interceptors, String path) {
+
+			mappings.add(new Mapping(wsHandler, path, handshakeHandler, interceptors));
+		}
+	}
+
+
+	private static class Mapping {
+
+		private final WebSocketHandler webSocketHandler;
+
+		private final String path;
+
+		private final HandshakeHandler handshakeHandler;
+
+		private final HandshakeInterceptor[] interceptors;
+
+		private final DefaultSockJsService sockJsService;
+
+
+		public Mapping(WebSocketHandler handler, String path, SockJsService sockJsService) {
+			this.webSocketHandler = handler;
+			this.path = path;
+			this.handshakeHandler = null;
+			this.interceptors = null;
+			this.sockJsService = (DefaultSockJsService) sockJsService;
+		}
+
+		public Mapping(WebSocketHandler h, String path, HandshakeHandler hh, HandshakeInterceptor[] interceptors) {
+			this.webSocketHandler = h;
+			this.path = path;
+			this.handshakeHandler = hh;
+			this.interceptors = interceptors;
+			this.sockJsService = null;
+		}
+	}
+
+}

spring-websocket/src/test/java/org/springframework/web/socket/server/config/WebSocketConfigurationTests.java

@@ -27,7 +27,6 @@ import org.junit.runners.Parameterized.Parameters;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
 import org.springframework.web.socket.AbstractWebSocketIntegrationTests;
 import org.springframework.web.socket.JettyTestServer;
 import org.springframework.web.socket.WebSocketSession;
@@ -54,33 +53,26 @@ public class WebSocketConfigurationTests extends AbstractWebSocketIntegrationTes
 	};
 
 
+	@Override
+	protected Class<?>[] getAnnotatedConfigClasses() {
+		return new Class<?>[] { TestWebSocketConfigurer.class };
+	}
+
 	@Test
 	public void registerWebSocketHandler() throws Exception {
 
-		AnnotationConfigWebApplicationContext cxt = new AnnotationConfigWebApplicationContext();
-		cxt.register(TestWebSocketConfigurer.class, getUpgradeStrategyConfigClass());
-
-		this.server.init(cxt);
-		this.server.start();
-
 		this.webSocketClient.doHandshake(new WebSocketHandlerAdapter(), getWsBaseUrl() + "/ws");
 
-		TestWebSocketHandler serverHandler = cxt.getBean(TestWebSocketHandler.class);
+		TestWebSocketHandler serverHandler = this.wac.getBean(TestWebSocketHandler.class);
 		assertTrue(serverHandler.latch.await(2, TimeUnit.SECONDS));
 	}
 
 	@Test
 	public void registerWebSocketHandlerWithSockJS() throws Exception {
 
-		AnnotationConfigWebApplicationContext cxt = new AnnotationConfigWebApplicationContext();
-		cxt.register(TestWebSocketConfigurer.class, getUpgradeStrategyConfigClass());
-
-		this.server.init(cxt);
-		this.server.start();
-
 		this.webSocketClient.doHandshake(new WebSocketHandlerAdapter(), getWsBaseUrl() + "/sockjs/websocket");
 
-		TestWebSocketHandler serverHandler = cxt.getBean(TestWebSocketHandler.class);
+		TestWebSocketHandler serverHandler = this.wac.getBean(TestWebSocketHandler.class);
 		assertTrue(serverHandler.latch.await(2, TimeUnit.SECONDS));
 	}