diff --git a/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/IntegrationWebSocketContainer.java b/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/IntegrationWebSocketContainer.java index 8e392628cf..5fbf0e14f6 100644 --- a/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/IntegrationWebSocketContainer.java +++ b/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/IntegrationWebSocketContainer.java @@ -1,5 +1,5 @@ /* - * Copyright 2014 the original author or authors. + * Copyright 2014-2015 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. @@ -181,10 +181,7 @@ public abstract class IntegrationWebSocketContainer implements DisposableBean { @Override public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception { - WebSocketSession removed = IntegrationWebSocketContainer.this.sessions.remove(session.getId()); - if (removed != null) { - IntegrationWebSocketContainer.this.sessions.remove(session.getId()); - } + IntegrationWebSocketContainer.this.sessions.remove(session.getId()); throw new Exception(exception); } diff --git a/spring-integration-websocket/src/test/java/org/springframework/integration/websocket/client/StompIntegrationTests.java b/spring-integration-websocket/src/test/java/org/springframework/integration/websocket/client/StompIntegrationTests.java index 42174d94ab..d688413d38 100644 --- a/spring-integration-websocket/src/test/java/org/springframework/integration/websocket/client/StompIntegrationTests.java +++ b/spring-integration-websocket/src/test/java/org/springframework/integration/websocket/client/StompIntegrationTests.java @@ -27,6 +27,7 @@ import java.lang.annotation.Retention; import java.lang.annotation.RetentionPolicy; import java.lang.annotation.Target; import java.nio.ByteBuffer; +import java.util.Collections; import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -84,6 +85,9 @@ import org.springframework.stereotype.Controller; import org.springframework.test.annotation.DirtiesContext; import org.springframework.test.context.ContextConfiguration; import org.springframework.test.context.junit4.SpringJUnit4ClassRunner; +import org.springframework.util.MultiValueMap; +import org.springframework.web.socket.WebSocketSession; +import org.springframework.web.socket.client.WebSocketClient; import org.springframework.web.socket.client.standard.StandardWebSocketClient; import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer; import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker; @@ -95,6 +99,9 @@ import org.springframework.web.socket.messaging.StompSubProtocolHandler; import org.springframework.web.socket.messaging.SubProtocolHandler; import org.springframework.web.socket.server.standard.TomcatRequestUpgradeStrategy; import org.springframework.web.socket.server.support.DefaultHandshakeHandler; +import org.springframework.web.socket.sockjs.client.SockJsClient; +import org.springframework.web.socket.sockjs.client.Transport; +import org.springframework.web.socket.sockjs.client.WebSocketTransport; /** * @author Artem Bilan @@ -108,6 +115,9 @@ public class StompIntegrationTests { @Value("#{server.serverContext}") private ApplicationContext serverContext; + @Autowired + private IntegrationWebSocketContainer clientWebSocketContainer; + @Autowired @Qualifier("webSocketOutputChannel") private MessageChannel webSocketOutputChannel; @@ -166,7 +176,7 @@ public class StompIntegrationTests { assertEquals(StompCommand.RECEIPT, headers.getCommand()); assertEquals("myReceipt", headers.getReceiptId()); - waitForSubscribe("increment"); + waitForSubscribe("/topic/increment"); headers = StompHeaderAccessor.create(StompCommand.SEND); headers.setSubscriptionId("subs1"); @@ -197,7 +207,7 @@ public class StompIntegrationTests { this.webSocketOutputChannel.send(message); - waitForSubscribe("foo"); + waitForSubscribe("/topic/foo"); this.webSocketOutputChannel.send(message2); @@ -251,7 +261,7 @@ public class StompIntegrationTests { this.webSocketOutputChannel.send(message); - waitForSubscribe("error"); + waitForSubscribe("/queue/error-user" + this.clientWebSocketContainer.getSession(null).getId()); this.webSocketOutputChannel.send(message2); @@ -284,7 +294,7 @@ public class StompIntegrationTests { this.webSocketOutputChannel.send(message); - waitForSubscribe("answer"); + waitForSubscribe("/queue/answer-user" + this.clientWebSocketContainer.getSession(null).getId()); this.webSocketOutputChannel.send(message2); @@ -307,18 +317,13 @@ public class StompIntegrationTests { assertTrue("The subscription for the '" + destination + "' destination hasn't been registered", n < 100); } - @SuppressWarnings("rawtypes") private boolean containsDestination(String destination, SubscriptionRegistry subscriptionRegistry) { - Map sessions = TestUtils.getPropertyValue(subscriptionRegistry, "subscriptionRegistry.sessions", Map.class); - for (Object info : sessions.values()) { - Map subscriptions = TestUtils.getPropertyValue(info, "destinationLookup", Map.class); - for (Object dest : subscriptions.keySet()) { - if (((String) dest).contains(destination)) { - return true; - } - } - } - return false; + StompHeaderAccessor stompHeaderAccessor = StompHeaderAccessor.create(StompCommand.MESSAGE); + stompHeaderAccessor.setDestination(destination); + Message message = MessageBuilder.createMessage(new byte[0], stompHeaderAccessor.toMessageHeaders()); + Object sessions = TestUtils.getPropertyValue(subscriptionRegistry, "subscriptionRegistry.sessions"); + MultiValueMap subscriptions = subscriptionRegistry.findSubscriptions(message); + return !subscriptions.isEmpty(); } @Configuration @@ -330,9 +335,14 @@ public class StompIntegrationTests { return new TomcatWebSocketTestServer(ServerConfig.class); } + @Bean + public WebSocketClient webSocketClient() { + return new SockJsClient(Collections.singletonList(new WebSocketTransport(new StandardWebSocketClient()))); + } + @Bean public IntegrationWebSocketContainer clientWebSocketContainer() { - return new ClientWebSocketContainer(new StandardWebSocketClient(), server().getWsBaseUrl() + "/ws/websocket"); + return new ClientWebSocketContainer(webSocketClient(), server().getWsBaseUrl() + "/ws"); } @Bean