Make NIO tests more reliable

Author: acogoluegnes

src/test/java/com/rabbitmq/client/test/JavaNioTest.java

@@ -2,6 +2,8 @@
 
 import com.rabbitmq.client.*;
 import com.rabbitmq.client.impl.nio.NioParams;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
 
 import java.io.IOException;
@@ -14,6 +16,25 @@
  */
 public class JavaNioTest {
 
+    public static final String QUEUE = "nio.queue";
+
+    private Connection testConnection;
+
+    @Before
+    public void init() throws Exception {
+        ConnectionFactory connectionFactory = new ConnectionFactory();
+        connectionFactory.useNio();
+        testConnection = connectionFactory.newConnection();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        if (testConnection != null) {
+            testConnection.createChannel().queueDelete(QUEUE);
+            testConnection.close();
+        }
+    }
+
     @Test
     public void connection() throws Exception {
         CountDownLatch latch = new CountDownLatch(1);
@@ -101,6 +122,18 @@ public void nioLoopCleaning() throws Exception {
         }
     }
 
+    @Test public void messageSize() throws Exception {
+        for (int i = 0; i < 50; i++) {
+            sendAndVerifyMessage(testConnection, 76390);
+        }
+    }
+
+    private void sendAndVerifyMessage(Connection connection, int size) throws Exception {
+        CountDownLatch latch = new CountDownLatch(1);
+        boolean messageReceived = basicGetBasicConsume(connection, QUEUE, latch, size);
+        assertTrue("Message has not been received", messageReceived);
+    }
+
     private Connection basicGetBasicConsume(ConnectionFactory connectionFactory, String queue, final CountDownLatch latch)
         throws IOException, TimeoutException {
         Connection connection = connectionFactory.newConnection();
@@ -121,6 +154,28 @@ public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProp
         return connection;
     }
 
+    private boolean basicGetBasicConsume(Connection connection, String queue, final CountDownLatch latch, int msgSize)
+        throws Exception {
+        Channel channel = connection.createChannel();
+        channel.queueDeclare(queue, false, false, false, null);
+        channel.queuePurge(queue);
+
+        channel.basicPublish("", queue, null, new byte[msgSize]);
+
+        final String tag = channel.basicConsume(queue, false, new DefaultConsumer(channel) {
+
+            @Override
+            public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProperties properties, byte[] body) throws IOException {
+                getChannel().basicAck(envelope.getDeliveryTag(), false);
+                latch.countDown();
+            }
+        });
+
+        boolean done = latch.await(20, TimeUnit.SECONDS);
+        channel.basicCancel(tag);
+        return done;
+    }
+
     private void safeClose(Connection connection) {
         if (connection != null) {
             try {

src/test/java/com/rabbitmq/client/test/ssl/NioTlsUnverifiedConnection.java

@@ -71,7 +71,7 @@ protected void releaseResources() throws IOException {
     @Test
     public void connectionGetConsume() throws Exception {
         CountDownLatch latch = new CountDownLatch(1);
-        connection = basicGetBasicConsume(connection, QUEUE, latch, 100 * 1000);
+        basicGetBasicConsume(connection, QUEUE, latch, 100 * 1000);
         boolean messagesReceived = latch.await(5, TimeUnit.SECONDS);
         assertTrue("Message has not been received", messagesReceived);
     }
@@ -120,30 +120,32 @@ public void configure(SSLEngine sslEngine) throws IOException {
 
     private void sendAndVerifyMessage(int size) throws Exception {
         CountDownLatch latch = new CountDownLatch(1);
-        connection = basicGetBasicConsume(connection, QUEUE, latch, size);
-        boolean messagesReceived = latch.await(20, TimeUnit.SECONDS);
-        assertTrue("Message has not been received", messagesReceived);
+        boolean messageReceived = basicGetBasicConsume(connection, QUEUE, latch, size);
+        assertTrue("Message has not been received", messageReceived);
     }
 
-    private Connection basicGetBasicConsume(Connection connection, String queue, final CountDownLatch latch, int msgSize)
-        throws IOException, TimeoutException {
+    private boolean basicGetBasicConsume(Connection connection, String queue, final CountDownLatch latch, int msgSize)
+        throws Exception {
         Channel channel = connection.createChannel();
         channel.queueDeclare(queue, false, false, false, null);
         channel.queuePurge(queue);
 
         channel.basicPublish("", queue, null, new byte[msgSize]);
 
-        channel.basicConsume(queue, false, new DefaultConsumer(channel) {
+        String tag = channel.basicConsume(queue, false, new DefaultConsumer(channel) {
 
             @Override
             public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProperties properties, byte[] body) throws IOException {
                 getChannel().basicAck(envelope.getDeliveryTag(), false);
                 latch.countDown();
-                getChannel().basicCancel(consumerTag);
             }
         });
 
-        return connection;
+        boolean messageReceived = latch.await(20, TimeUnit.SECONDS);
+
+        channel.basicCancel(tag);
+
+        return messageReceived;
     }
 
 }