Polish WebSocket samples

spring-boot-samples/spring-boot-sample-websocket-undertow/src/test/java/samples/websocket/echo/CustomContainerWebSocketsApplicationTests.java

@@ -42,10 +42,10 @@ import org.springframework.util.SocketUtils;
 import org.springframework.web.socket.client.WebSocketConnectionManager;
 import org.springframework.web.socket.client.standard.StandardWebSocketClient;
 
+import samples.websocket.SampleUndertowWebSocketsApplication;
 import samples.websocket.client.GreetingService;
 import samples.websocket.client.SimpleClientWebSocketHandler;
 import samples.websocket.client.SimpleGreetingService;
-import samples.websocket.config.SampleUndertowWebSocketsApplication;
 import samples.websocket.echo.CustomContainerWebSocketsApplicationTests.CustomContainerConfiguration;
 
 import static org.junit.Assert.assertEquals;
@@ -63,14 +63,6 @@ public class CustomContainerWebSocketsApplicationTests {
 
 	private static int PORT = SocketUtils.findAvailableTcpPort();
 
-	@Configuration
-	protected static class CustomContainerConfiguration {
-		@Bean
-		public EmbeddedServletContainerFactory embeddedServletContainerFactory() {
-			return new UndertowEmbeddedServletContainerFactory("/ws", PORT);
-		}
-	}
-
 	@Test
 	public void echoEndpoint() throws Exception {
 		ConfigurableApplicationContext context = new SpringApplicationBuilder(
@@ -99,6 +91,16 @@ public class CustomContainerWebSocketsApplicationTests {
 		assertEquals("Reversed: !dlrow olleH", messagePayloadReference.get());
 	}
 
+	@Configuration
+	protected static class CustomContainerConfiguration {
+
+		@Bean
+		public EmbeddedServletContainerFactory embeddedServletContainerFactory() {
+			return new UndertowEmbeddedServletContainerFactory("/ws", PORT);
+		}
+
+	}
+
 	@Configuration
 	static class ClientConfiguration implements CommandLineRunner {
 
@@ -145,6 +147,7 @@ public class CustomContainerWebSocketsApplicationTests {
 		public GreetingService greetingService() {
 			return new SimpleGreetingService();
 		}
+
 	}
 
 }