Refactor test to use less memory

Author: markt-asf

test/org/apache/tomcat/websocket/TestWebSocketFrameClient.java

@@ -77,7 +77,7 @@ public void beforeRequest(Map<String, List<String>> headers) {
         Session wsSession = wsContainer.connectToServer(TesterProgrammaticEndpoint.class, clientEndpointConfig,
                 new URI("ws://localhost:" + getPort() + TesterFirehoseServer.PATH));
         CountDownLatch latch = new CountDownLatch(TesterFirehoseServer.MESSAGE_COUNT);
-        BasicText handler = new BasicText(latch);
+        BasicText handler = new BasicText(latch, TesterFirehoseServer.MESSAGE);
         wsSession.addMessageHandler(handler);
         wsSession.getBasicRemote().sendText("Hello");
 
@@ -87,11 +87,7 @@ public void beforeRequest(Map<String, List<String>> headers) {
         // if the right number of messages arrived
         handler.getLatch().await(TesterFirehoseServer.WAIT_TIME_MILLIS, TimeUnit.MILLISECONDS);
 
-        Queue<String> messages = handler.getMessages();
-        Assert.assertEquals(TesterFirehoseServer.MESSAGE_COUNT, messages.size());
-        for (String message : messages) {
-            Assert.assertEquals(TesterFirehoseServer.MESSAGE, message);
-        }
+        Assert.assertEquals(TesterFirehoseServer.MESSAGE_COUNT, handler.getMessageCount());
     }
 
     @Test

test/org/apache/tomcat/websocket/TesterMessageCountClient.java

@@ -22,6 +22,7 @@
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import javax.websocket.ClientEndpoint;
 import javax.websocket.CloseReason;
@@ -144,6 +145,7 @@ public void onMessage(ByteBuffer message) {
     public static class BasicText extends BasicHandler<String> {
 
         private final String expected;
+        private final AtomicInteger messageCount;
 
         public BasicText(CountDownLatch latch) {
             this(latch, null);
@@ -152,6 +154,11 @@ public BasicText(CountDownLatch latch) {
         public BasicText(CountDownLatch latch, String expected) {
             super(latch);
             this.expected = expected;
+            if (expected == null) {
+                messageCount = null;
+            } else {
+                messageCount = new AtomicInteger(0);
+            }
         }
 
         @Override
@@ -162,11 +169,20 @@ public void onMessage(String message) {
                 if (!expected.equals(message)) {
                     throw new IllegalStateException("Expected: [" + expected + "]\r\n" + "Was:      [" + message + "]");
                 }
+                messageCount.incrementAndGet();
             }
             if (getLatch() != null) {
                 getLatch().countDown();
             }
         }
+
+        public int getMessageCount() {
+            if (expected == null) {
+                return getMessages().size();
+            } else {
+                return messageCount.get();
+            }
+        }
     }
 
     public static class SleepingText implements MessageHandler.Whole<String> {