From 05cd1694c21f06450bb0cb103caf1e345d5b1389 Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@vmware.com>
Date: Wed, 20 Feb 2013 10:41:24 -0500
Subject: [PATCH] INT-2940 Fix Race Condition in TCP Connection Test

Assertion that server side connection was established could
occur after the connection was established, but before it was
added to the collection.

Add a latch.
---
 .../tcp/connection/ConnectionFactoryTests.java | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/spring-integration-ip/src/test/java/org/springframework/integration/ip/tcp/connection/ConnectionFactoryTests.java b/spring-integration-ip/src/test/java/org/springframework/integration/ip/tcp/connection/ConnectionFactoryTests.java
index 5d29394d81..c5bc66c0f0 100644
--- a/spring-integration-ip/src/test/java/org/springframework/integration/ip/tcp/connection/ConnectionFactoryTests.java
+++ b/spring-integration-ip/src/test/java/org/springframework/integration/ip/tcp/connection/ConnectionFactoryTests.java
@@ -18,13 +18,20 @@ package org.springframework.integration.ip.tcp.connection;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 import org.junit.Test;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import org.springframework.context.ApplicationEvent;
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.integration.Message;
@@ -35,6 +42,7 @@ import org.springframework.integration.test.util.SocketUtils;
 
 /**
  * @author Gary Russell
+ * @author Artem Bilan
  * @since 3.0
  *
  */
@@ -53,6 +61,15 @@ public class ConnectionFactoryTests {
 		TcpNetServerConnectionFactory serverFactory = new TcpNetServerConnectionFactory(port);
 		serverFactory.setBeanName("serverFactory");
 		serverFactory.setApplicationEventPublisher(publisher);
+		serverFactory = spy(serverFactory);
+		final CountDownLatch serverConnectionInitLatch = new CountDownLatch(1);
+		doAnswer(new Answer<Object>() {
+			public Object answer(InvocationOnMock invocation) throws Throwable {
+				Object result = invocation.callRealMethod();
+				serverConnectionInitLatch.countDown();
+				return result;
+			}
+		}).when(serverFactory).wrapConnection(any(TcpConnectionSupport.class));
 		TcpReceivingChannelAdapter adapter = new TcpReceivingChannelAdapter();
 		adapter.setConnectionFactory(serverFactory);
 		adapter.start();
@@ -70,6 +87,7 @@ public class ConnectionFactoryTests {
 		List<String> clients = clientFactory.getOpenConnectionIds();
 		assertEquals(1, clients.size());
 		assertTrue(clients.contains(client.getConnectionId()));
+		assertTrue("Server connection failed to register", serverConnectionInitLatch.await(1, TimeUnit.SECONDS));
 		List<String> servers = serverFactory.getOpenConnectionIds();
 		assertEquals(1, servers.size());
 		assertTrue(serverFactory.closeConnection(servers.get(0)));