ServerEndpointExporter can initialize itself based on a late-provided ServletContext as well (for Boot)

Also allows for direct setting of a ServerContainer and for custom triggering of endpoint registration.

Issue: SPR-12109

spring-websocket/src/test/java/org/springframework/web/socket/server/standard/ServerEndpointExporterTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -16,6 +16,7 @@
 
 package org.springframework.web.socket.server.standard;
 
+import javax.servlet.ServletContext;
 import javax.websocket.Endpoint;
 import javax.websocket.EndpointConfig;
 import javax.websocket.Session;
@@ -24,6 +25,7 @@ import javax.websocket.server.ServerEndpoint;
 
 import org.junit.Before;
 import org.junit.Test;
+
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.mock.web.test.MockServletContext;
@@ -35,37 +37,59 @@ import static org.mockito.Mockito.*;
  * Test fixture for {@link ServerEndpointExporter}.
  *
  * @author Rossen Stoyanchev
+ * @author Juergen Hoeller
  */
 public class ServerEndpointExporterTests {
 
 	private ServerContainer serverContainer;
 
-	private ServerEndpointExporter exporter;
+	private ServletContext servletContext;
 
 	private AnnotationConfigWebApplicationContext webAppContext;
 
+	private ServerEndpointExporter exporter;
+
 
 	@Before
 	public void setup() {
 		this.serverContainer = mock(ServerContainer.class);
 
-		MockServletContext servletContext = new MockServletContext();
-		servletContext.setAttribute("javax.websocket.server.ServerContainer", this.serverContainer);
+		this.servletContext = new MockServletContext();
+		this.servletContext.setAttribute("javax.websocket.server.ServerContainer", this.serverContainer);
 
 		this.webAppContext = new AnnotationConfigWebApplicationContext();
 		this.webAppContext.register(Config.class);
-		this.webAppContext.setServletContext(servletContext);
+		this.webAppContext.setServletContext(this.servletContext);
 		this.webAppContext.refresh();
 
 		this.exporter = new ServerEndpointExporter();
-		this.exporter.setApplicationContext(this.webAppContext);
 	}
 
 
 	@Test
-	public void addAnnotatedEndpointBean() throws Exception {
-
+	public void addAnnotatedEndpointBeans() throws Exception {
 		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class);
+		this.exporter.setApplicationContext(this.webAppContext);
+		this.exporter.afterPropertiesSet();
+
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpoint.class);
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpointBean.class);
+	}
+
+	@Test
+	public void addAnnotatedEndpointBeansWithServletContextOnly() throws Exception {
+		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class, AnnotatedDummyEndpointBean.class);
+		this.exporter.setServletContext(this.servletContext);
+		this.exporter.afterPropertiesSet();
+
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpoint.class);
+		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpointBean.class);
+	}
+
+	@Test
+	public void addAnnotatedEndpointBeansWithServerContainerOnly() throws Exception {
+		this.exporter.setAnnotatedEndpointClasses(AnnotatedDummyEndpoint.class, AnnotatedDummyEndpointBean.class);
+		this.exporter.setServerContainer(this.serverContainer);
 		this.exporter.afterPropertiesSet();
 
 		verify(this.serverContainer).addEndpoint(AnnotatedDummyEndpoint.class);
@@ -74,10 +98,31 @@ public class ServerEndpointExporterTests {
 
 	@Test
 	public void addServerEndpointConfigBean() throws Exception {
+		this.exporter.setApplicationContext(this.webAppContext);
+		this.exporter.afterPropertiesSet();
 
 		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
 		this.exporter.postProcessAfterInitialization(endpointRegistration, "dummyEndpoint");
+		verify(this.serverContainer).addEndpoint(endpointRegistration);
+	}
 
+	@Test
+	public void addServerEndpointConfigBeanWithServletContextOnly() throws Exception {
+		this.exporter.setServletContext(this.servletContext);
+		this.exporter.afterPropertiesSet();
+
+		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
+		this.exporter.postProcessAfterInitialization(endpointRegistration, "dummyEndpoint");
+		verify(this.serverContainer).addEndpoint(endpointRegistration);
+	}
+
+	@Test
+	public void addServerEndpointConfigBeanWithServerContainerOnly() throws Exception {
+		this.exporter.setServerContainer(this.serverContainer);
+		this.exporter.afterPropertiesSet();
+
+		ServerEndpointRegistration endpointRegistration = new ServerEndpointRegistration("/dummy", new DummyEndpoint());
+		this.exporter.postProcessAfterInitialization(endpointRegistration, "dummyEndpoint");
 		verify(this.serverContainer).addEndpoint(endpointRegistration);
 	}
 
@@ -89,14 +134,17 @@ public class ServerEndpointExporterTests {
 		}
 	}
 
+
 	@ServerEndpoint("/path")
 	private static class AnnotatedDummyEndpoint {
 	}
 
+
 	@ServerEndpoint("/path")
 	private static class AnnotatedDummyEndpointBean {
 	}
 
+
 	@Configuration
 	static class Config {