From 9b1aab070350c40e8028b385f3b3805fbf4a40c7 Mon Sep 17 00:00:00 2001
From: Artem Bilan <abilan@pivotal.io>
Date: Wed, 27 May 2015 14:23:48 -0400
Subject: [PATCH] Fix STOMP subscription race condition

Even if we receive a SUBSCRIBE event, it doesn't that the subscription has been registered.
 So, fix the test to wait for the real subscription.
---
 ...annelAdapterWebSocketIntegrationTests.java | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/spring-integration-stomp/src/test/java/org/springframework/integration/stomp/inbound/StompInboundChannelAdapterWebSocketIntegrationTests.java b/spring-integration-stomp/src/test/java/org/springframework/integration/stomp/inbound/StompInboundChannelAdapterWebSocketIntegrationTests.java
index 2408df3a8f..ccd18c6e3a 100644
--- a/spring-integration-stomp/src/test/java/org/springframework/integration/stomp/inbound/StompInboundChannelAdapterWebSocketIntegrationTests.java
+++ b/spring-integration-stomp/src/test/java/org/springframework/integration/stomp/inbound/StompInboundChannelAdapterWebSocketIntegrationTests.java
@@ -22,6 +22,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
 
 import java.util.Collections;
 import java.util.Map;
@@ -45,6 +46,7 @@ import org.springframework.integration.stomp.StompSessionManager;
 import org.springframework.integration.stomp.WebSocketStompSessionManager;
 import org.springframework.integration.stomp.event.StompIntegrationEvent;
 import org.springframework.integration.stomp.event.StompReceiptEvent;
+import org.springframework.integration.test.util.TestUtils;
 import org.springframework.integration.websocket.TomcatWebSocketTestServer;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHandlingException;
@@ -52,6 +54,8 @@ import org.springframework.messaging.PollableChannel;
 import org.springframework.messaging.converter.MappingJackson2MessageConverter;
 import org.springframework.messaging.converter.MessageConversionException;
 import org.springframework.messaging.simp.SimpMessagingTemplate;
+import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
+import org.springframework.messaging.simp.broker.SubscriptionRegistry;
 import org.springframework.messaging.simp.config.MessageBrokerRegistry;
 import org.springframework.messaging.simp.stomp.StompCommand;
 import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
@@ -111,6 +115,8 @@ public class StompInboundChannelAdapterWebSocketIntegrationTests {
 		assertEquals(StompCommand.SUBSCRIBE, stompReceiptEvent.getStompCommand());
 		assertEquals("/topic/myTopic", stompReceiptEvent.getDestination());
 
+		waitForSubscribe("myTopic");
+
 		SimpMessagingTemplate messagingTemplate = this.serverContext.getBean("brokerMessagingTemplate",
 				SimpMessagingTemplate.class);
 
@@ -140,6 +146,9 @@ public class StompInboundChannelAdapterWebSocketIntegrationTests {
 		this.stompInboundChannelAdapter.addDestination("/topic/myTopic");
 		receive = this.stompEvents.receive(10000);
 		assertNotNull(receive);
+
+		waitForSubscribe("myTopic");
+
 		messagingTemplate.convertAndSend("/topic/myTopic", "foo");
 		receive = this.errorChannel.receive(10000);
 		assertNotNull(receive);
@@ -151,6 +160,35 @@ public class StompInboundChannelAdapterWebSocketIntegrationTests {
 		assertThat(throwable.getMessage(), containsString("No suitable converter, payloadType=interface java.util.Map"));
 	}
 
+	private void waitForSubscribe(String destination) throws InterruptedException {
+		SimpleBrokerMessageHandler serverBrokerMessageHandler =
+				this.serverContext.getBean("simpleBrokerMessageHandler", SimpleBrokerMessageHandler.class);
+
+		SubscriptionRegistry subscriptionRegistry = serverBrokerMessageHandler.getSubscriptionRegistry();
+
+		int n = 0;
+		while (!containsDestination(destination, subscriptionRegistry) && n++ < 100) {
+			Thread.sleep(100);
+		}
+
+		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;
+	}
+
+
 	// STOMP Client
 
 	@Configuration