Polish WebSocketMessageBrokerConfigurationSupportTests

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

 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -47,14 +47,12 @@ import org.springframework.messaging.support.ImmutableMessageChannelInterceptor;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.stereotype.Controller;
-import org.springframework.web.servlet.HandlerMapping;
 import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.config.WebSocketMessageBrokerStats;
 import org.springframework.web.socket.handler.TestWebSocketSession;
 import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
-import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
 import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.messaging.StompTextMessageBuilder;
 import org.springframework.web.socket.messaging.SubProtocolHandler;
@@ -65,29 +63,29 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 
 /**
- * Test fixture for {@link WebSocketMessageBrokerConfigurationSupport}.
+ * Tests for {@link WebSocketMessageBrokerConfigurationSupport}.
  *
  * @author Rossen Stoyanchev
  * @author Sebastien Deleuze
  */
-public class WebSocketMessageBrokerConfigurationSupportTests {
+class WebSocketMessageBrokerConfigurationSupportTests {
 
 	@Test
-	public void handlerMapping() {
-		ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
-		SimpleUrlHandlerMapping hm = (SimpleUrlHandlerMapping) config.getBean(HandlerMapping.class);
+	void handlerMapping() {
+		ApplicationContext context = createContext(TestChannelConfig.class, TestConfigurer.class);
+		SimpleUrlHandlerMapping hm = context.getBean(SimpleUrlHandlerMapping.class);
 		assertThat(hm.getOrder()).isEqualTo(1);
 
 		Map<String, Object> handlerMap = hm.getHandlerMap();
-		assertThat(handlerMap.size()).isEqualTo(1);
+		assertThat(handlerMap).hasSize(1);
 		assertThat(handlerMap.get("/simpleBroker")).isNotNull();
 	}
 
 	@Test
-	public void clientInboundChannelSendMessage() throws Exception {
-		ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
-		TestChannel channel = config.getBean("clientInboundChannel", TestChannel.class);
-		SubProtocolWebSocketHandler webSocketHandler = config.getBean(SubProtocolWebSocketHandler.class);
+	void clientInboundChannelSendMessage() throws Exception {
+		ApplicationContext context = createContext(TestChannelConfig.class, TestConfigurer.class);
+		TestChannel channel = context.getBean("clientInboundChannel", TestChannel.class);
+		SubProtocolWebSocketHandler webSocketHandler = context.getBean(SubProtocolWebSocketHandler.class);
 
 		List<ChannelInterceptor> interceptors = channel.getInterceptors();
 		assertThat(interceptors.get(interceptors.size() - 1).getClass()).isEqualTo(ImmutableMessageChannelInterceptor.class);
@@ -108,66 +106,66 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
 	}
 
 	@Test
-	public void clientOutboundChannel() {
-		ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
-		TestChannel channel = config.getBean("clientOutboundChannel", TestChannel.class);
+	void clientOutboundChannel() {
+		ApplicationContext context = createContext(TestChannelConfig.class, TestConfigurer.class);
+		TestChannel channel = context.getBean("clientOutboundChannel", TestChannel.class);
 		Set<MessageHandler> handlers = channel.getSubscribers();
 
 		List<ChannelInterceptor> interceptors = channel.getInterceptors();
 		assertThat(interceptors.get(interceptors.size() - 1).getClass()).isEqualTo(ImmutableMessageChannelInterceptor.class);
 
-		assertThat(handlers.size()).isEqualTo(1);
-		assertThat(handlers.contains(config.getBean(SubProtocolWebSocketHandler.class))).isTrue();
+		assertThat(handlers).hasSize(1);
+		assertThat(handlers).contains(context.getBean(SubProtocolWebSocketHandler.class));
 	}
 
 	@Test
-	public void brokerChannel() {
-		ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
-		TestChannel channel = config.getBean("brokerChannel", TestChannel.class);
+	void brokerChannel() {
+		ApplicationContext context = createContext(TestChannelConfig.class, TestConfigurer.class);
+		TestChannel channel = context.getBean("brokerChannel", TestChannel.class);
 		Set<MessageHandler> handlers = channel.getSubscribers();
 
 		List<ChannelInterceptor> interceptors = channel.getInterceptors();
 		assertThat(interceptors.get(interceptors.size() - 1).getClass()).isEqualTo(ImmutableMessageChannelInterceptor.class);
 
-		assertThat(handlers.size()).isEqualTo(2);
-		assertThat(handlers.contains(config.getBean(SimpleBrokerMessageHandler.class))).isTrue();
-		assertThat(handlers.contains(config.getBean(UserDestinationMessageHandler.class))).isTrue();
+		assertThat(handlers).hasSize(2);
+		assertThat(handlers).contains(context.getBean(SimpleBrokerMessageHandler.class));
+		assertThat(handlers).contains(context.getBean(UserDestinationMessageHandler.class));
 	}
 
 	@Test
-	public void webSocketHandler() {
-		ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
-		SubProtocolWebSocketHandler subWsHandler = config.getBean(SubProtocolWebSocketHandler.class);
+	void webSocketHandler() {
+		ApplicationContext context = createContext(TestChannelConfig.class, TestConfigurer.class);
+		SubProtocolWebSocketHandler subWsHandler = context.getBean(SubProtocolWebSocketHandler.class);
 
-		assertThat(subWsHandler.getSendBufferSizeLimit()).isEqualTo((1024 * 1024));
-		assertThat(subWsHandler.getSendTimeLimit()).isEqualTo((25 * 1000));
-		assertThat(subWsHandler.getTimeToFirstMessage()).isEqualTo((30 * 1000));
+		assertThat(subWsHandler.getSendBufferSizeLimit()).isEqualTo(1024 * 1024);
+		assertThat(subWsHandler.getSendTimeLimit()).isEqualTo(25 * 1000);
+		assertThat(subWsHandler.getTimeToFirstMessage()).isEqualTo(30 * 1000);
 
 		Map<String, SubProtocolHandler> handlerMap = subWsHandler.getProtocolHandlerMap();
 		StompSubProtocolHandler protocolHandler = (StompSubProtocolHandler) handlerMap.get("v12.stomp");
-		assertThat(protocolHandler.getMessageSizeLimit()).isEqualTo((128 * 1024));
+		assertThat(protocolHandler.getMessageSizeLimit()).isEqualTo(128 * 1024);
 	}
 
 	@Test
-	public void taskScheduler() {
-		ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
+	void taskScheduler() {
+		ApplicationContext context = createContext(TestChannelConfig.class, TestConfigurer.class);
 
 		String name = "messageBrokerSockJsTaskScheduler";
-		ThreadPoolTaskScheduler taskScheduler = config.getBean(name, ThreadPoolTaskScheduler.class);
+		ThreadPoolTaskScheduler taskScheduler = context.getBean(name, ThreadPoolTaskScheduler.class);
 		ScheduledThreadPoolExecutor executor = taskScheduler.getScheduledThreadPoolExecutor();
 		assertThat(executor.getCorePoolSize()).isEqualTo(Runtime.getRuntime().availableProcessors());
 		assertThat(executor.getRemoveOnCancelPolicy()).isTrue();
 
-		SimpleBrokerMessageHandler handler = config.getBean(SimpleBrokerMessageHandler.class);
+		SimpleBrokerMessageHandler handler = context.getBean(SimpleBrokerMessageHandler.class);
 		assertThat(handler.getTaskScheduler()).isNotNull();
-		assertThat(handler.getHeartbeatValue()).isEqualTo(new long[] {15000, 15000});
+		assertThat(handler.getHeartbeatValue()).containsExactly(15000, 15000);
 	}
 
 	@Test
-	public void webSocketMessageBrokerStats() {
-		ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
+	void webSocketMessageBrokerStats() {
+		ApplicationContext context = createContext(TestChannelConfig.class, TestConfigurer.class);
 		String name = "webSocketMessageBrokerStats";
-		WebSocketMessageBrokerStats stats = config.getBean(name, WebSocketMessageBrokerStats.class);
+		WebSocketMessageBrokerStats stats = context.getBean(name, WebSocketMessageBrokerStats.class);
 		String actual = stats.toString();
 		String expected = "WebSocketSession\\[0 current WS\\(0\\)-HttpStream\\(0\\)-HttpPoll\\(0\\), " +
 				"0 total, 0 closed abnormally \\(0 connect failure, 0 send limit, 0 transport error\\)], " +
@@ -177,16 +175,16 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
 				"outboundChannel\\[pool size = \\d, active threads = \\d, queued tasks = \\d, completed tasks = \\d], " +
 				"sockJsScheduler\\[pool size = \\d, active threads = \\d, queued tasks = \\d, completed tasks = \\d]";
 
-		assertThat(actual.matches(expected)).as("\nExpected: " + expected.replace("\\", "") + "\n  Actual: " + actual).isTrue();
+		assertThat(actual).matches(expected);
 	}
 
 	@Test
-	public void webSocketHandlerDecorator() throws Exception {
-		ApplicationContext config = createConfig(WebSocketHandlerDecoratorConfig.class);
-		WebSocketHandler handler = config.getBean(SubProtocolWebSocketHandler.class);
+	void webSocketHandlerDecorator() throws Exception {
+		ApplicationContext context = createContext(WebSocketHandlerDecoratorConfig.class);
+		WebSocketHandler handler = context.getBean(SubProtocolWebSocketHandler.class);
 		assertThat(handler).isNotNull();
 
-		SimpleUrlHandlerMapping mapping = (SimpleUrlHandlerMapping) config.getBean("stompWebSocketHandlerMapping");
+		SimpleUrlHandlerMapping mapping = context.getBean("stompWebSocketHandlerMapping", SimpleUrlHandlerMapping.class);
 		WebSocketHttpRequestHandler httpHandler = (WebSocketHttpRequestHandler) mapping.getHandlerMap().get("/test");
 		handler = httpHandler.getWebSocketHandler();
 
@@ -196,11 +194,8 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
 	}
 
 
-	private ApplicationContext createConfig(Class<?>... configClasses) {
-		AnnotationConfigApplicationContext config = new AnnotationConfigApplicationContext();
-		config.register(configClasses);
-		config.refresh();
-		return config;
+	private ApplicationContext createContext(Class<?>... configClasses) {
+		return new AnnotationConfigApplicationContext(configClasses);
 	}
 
 
@@ -286,17 +281,12 @@ public class WebSocketMessageBrokerConfigurationSupportTests {
 
 		@Override
 		protected void configureWebSocketTransport(WebSocketTransportRegistration registry) {
-			registry.addDecoratorFactory(new WebSocketHandlerDecoratorFactory() {
-				@Override
-				public WebSocketHandlerDecorator decorate(WebSocketHandler handler) {
-					return new WebSocketHandlerDecorator(handler) {
-						@Override
-						public void afterConnectionEstablished(WebSocketSession session) throws Exception {
-							session.getAttributes().put("decorated", true);
-							super.afterConnectionEstablished(session);
-						}
-					};
-				}
+			registry.addDecoratorFactory(handler -> new WebSocketHandlerDecorator(handler) {
+					@Override
+					public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+						session.getAttributes().put("decorated", true);
+						super.afterConnectionEstablished(session);
+					}
 			});
 		}
 	}