Adding close validation to TestABCase1

4 files changed, 83 insertions, 10 deletions

diff --git a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/AbstractABCase.java b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/AbstractABCase.java
index ead17962a4..0723b7155a 100644
--- a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/AbstractABCase.java
+++ b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/AbstractABCase.java
@@ -2,6 +2,10 @@ package org.eclipse.jetty.websocket.server.ab;
 
 import java.nio.ByteBuffer;
 
+import org.eclipse.jetty.io.ByteBufferPool;
+import org.eclipse.jetty.io.StandardByteBufferPool;
+import org.eclipse.jetty.websocket.api.WebSocketPolicy;
+import org.eclipse.jetty.websocket.protocol.Generator;
 import org.eclipse.jetty.websocket.server.SimpleServletServer;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -14,9 +18,20 @@ public abstract class AbstractABCase
     private static final byte[] MASK =
     { 0x12, 0x34, 0x56, 0x78 };
 
+    protected static Generator strictGenerator;
+    protected static Generator laxGenerator;
     protected static SimpleServletServer server;
 
     @BeforeClass
+    public static void initGenerators()
+    {
+        WebSocketPolicy policy = WebSocketPolicy.newClientPolicy();
+        ByteBufferPool bufferPool = new StandardByteBufferPool();
+        strictGenerator = new Generator(policy,bufferPool,true);
+        laxGenerator = new Generator(policy,bufferPool,false);
+    }
+
+    @BeforeClass
     public static void startServer() throws Exception
     {
         server = new SimpleServletServer(new ABServlet());
diff --git a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/TestABCase1.java b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/TestABCase1.java
index f6b6e0bb1b..cc3a853003 100644
--- a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/TestABCase1.java
+++ b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/ab/TestABCase1.java
@@ -23,6 +23,8 @@ import java.util.Queue;
 import java.util.concurrent.TimeUnit;
 
 import org.eclipse.jetty.util.BufferUtil;
+import org.eclipse.jetty.websocket.api.StatusCode;
+import org.eclipse.jetty.websocket.protocol.CloseInfo;
 import org.eclipse.jetty.websocket.protocol.Generator;
 import org.eclipse.jetty.websocket.protocol.OpCode;
 import org.eclipse.jetty.websocket.protocol.WebSocketFrame;
@@ -45,22 +47,40 @@ public class TestABCase1 extends AbstractABCase
             ByteBuffer buf = ByteBuffer.allocate(Generator.OVERHEAD);
             BufferUtil.clearToFill(buf);
 
+            // Prepare Frame
             buf.put((byte)(0x00 | FIN | opcode.getCode()));
             putPayloadLength(buf,0);
             putMask(buf);
 
+            // Write Data Frame
             BufferUtil.flipToFlush(buf,0);
             client.writeRaw(buf);
 
-            // Read frame
-            Queue<WebSocketFrame> frames = client.readFrames(1,TimeUnit.MILLISECONDS,500);
+            // Prepare Close Frame
+            CloseInfo close = new CloseInfo(StatusCode.NORMAL);
+            buf = strictGenerator.generate(close.asFrame());
+
+            // Write Close Frame
+            client.writeRaw(buf);
+            client.flush();
+
+            // Read frames
+            Queue<WebSocketFrame> frames = client.readFrames(2,TimeUnit.MILLISECONDS,500);
+
+            // Validate echo'd frame
             WebSocketFrame frame = frames.remove();
             Assert.assertThat("frame should be " + opcode + " frame",frame.getOpCode(),is(opcode));
             Assert.assertThat(opcode + ".payloadLength",frame.getPayloadLength(),is(0));
+
+            // Validate close
+            frame = frames.remove();
+            Assert.assertThat("CLOSE.frame.opcode",frame.getOpCode(),is(OpCode.CLOSE));
+            close = new CloseInfo(frame);
+            Assert.assertThat("CLOSE.statusCode",close.getStatusCode(),is(StatusCode.NORMAL));
         }
         finally
         {
-            client.close();
+            client.disconnect();
         }
     }
 
@@ -76,24 +96,42 @@ public class TestABCase1 extends AbstractABCase
             ByteBuffer buf = ByteBuffer.allocate(payload.length + Generator.OVERHEAD);
             BufferUtil.clearToFill(buf);
 
+            // Prepare Frame
             buf.put((byte)(0x00 | FIN | opcode.getCode()));
             putPayloadLength(buf,payload.length);
             putMask(buf);
             buf.put(masked(payload));
 
+            // Write Data Frame
             BufferUtil.flipToFlush(buf,0);
             client.writeRaw(buf);
 
-            // Read frame
-            Queue<WebSocketFrame> frames = client.readFrames(1,TimeUnit.MILLISECONDS,500);
+            // Prepare Close Frame
+            CloseInfo close = new CloseInfo(StatusCode.NORMAL);
+            buf = strictGenerator.generate(close.asFrame());
+
+            // Write Close Frame
+            client.writeRaw(buf);
+            client.flush();
+
+            // Read frames
+            Queue<WebSocketFrame> frames = client.readFrames(2,TimeUnit.MILLISECONDS,500);
+
+            // Validate echo'd frame
             WebSocketFrame frame = frames.remove();
             Assert.assertThat("frame should be " + opcode + " frame",frame.getOpCode(),is(opcode));
             Assert.assertThat(opcode + ".payloadLength",frame.getPayloadLength(),is(payload.length));
             ByteBufferAssert.assertEquals(opcode + ".payload",payload,frame.getPayload());
+
+            // Validate close
+            frame = frames.remove();
+            Assert.assertThat("CLOSE.frame.opcode",frame.getOpCode(),is(OpCode.CLOSE));
+            close = new CloseInfo(frame);
+            Assert.assertThat("CLOSE.statusCode",close.getStatusCode(),is(StatusCode.NORMAL));
         }
         finally
         {
-            client.close();
+            client.disconnect();
         }
     }
 
@@ -109,11 +147,13 @@ public class TestABCase1 extends AbstractABCase
             ByteBuffer buf = ByteBuffer.allocate(payload.length + Generator.OVERHEAD);
             BufferUtil.clearToFill(buf);
 
+            // Prepare Frame
             buf.put((byte)(0x00 | FIN | opcode.getCode()));
             putPayloadLength(buf,payload.length);
             putMask(buf);
             buf.put(masked(payload));
 
+            // Write frame, in small blocks of segmentSize
             BufferUtil.flipToFlush(buf,0);
             int origLimit = buf.limit();
             int limit = buf.limit();
@@ -132,16 +172,32 @@ public class TestABCase1 extends AbstractABCase
                 client.flush();
             }
 
-            // Read frame
-            Queue<WebSocketFrame> frames = client.readFrames(1,TimeUnit.MILLISECONDS,500);
+            // Prepare Close Frame
+            CloseInfo close = new CloseInfo(StatusCode.NORMAL);
+            buf = strictGenerator.generate(close.asFrame());
+
+            // Write Close Frame
+            client.writeRaw(buf);
+            client.flush();
+
+            // Read frames
+            Queue<WebSocketFrame> frames = client.readFrames(2,TimeUnit.MILLISECONDS,500);
+
+            // Validate echo'd frame
             WebSocketFrame frame = frames.remove();
             Assert.assertThat("frame should be " + opcode + " frame",frame.getOpCode(),is(opcode));
             Assert.assertThat(opcode + ".payloadLength",frame.getPayloadLength(),is(payload.length));
             ByteBufferAssert.assertEquals(opcode + ".payload",payload,frame.getPayload());
+
+            // Validate close
+            frame = frames.remove();
+            Assert.assertThat("CLOSE.frame.opcode",frame.getOpCode(),is(OpCode.CLOSE));
+            close = new CloseInfo(frame);
+            Assert.assertThat("CLOSE.statusCode",close.getStatusCode(),is(StatusCode.NORMAL));
         }
         finally
         {
-            client.close();
+            client.disconnect();
         }
     }
 
diff --git a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/blockhead/BlockheadClient.java b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/blockhead/BlockheadClient.java
index bebd28ff4c..65d09e5f36 100644
--- a/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/blockhead/BlockheadClient.java
+++ b/jetty-websocket/websocket-server/src/test/java/org/eclipse/jetty/websocket/server/blockhead/BlockheadClient.java
@@ -415,7 +415,7 @@ public class BlockheadClient implements IncomingFrames, OutgoingFrames
                 }
                 if (!debug && (System.currentTimeMillis() > expireOn))
                 {
-                    throw new TimeoutException(String.format("Timeout reading all [%d] expected frames. (managed to read [%d] frames)",expectedCount,
+                    throw new TimeoutException(String.format("Timeout reading all %d expected frames. (managed to only read %d frame(s))",expectedCount,
                             incomingFrameQueue.size()));
                 }
             }
diff --git a/jetty-websocket/websocket-server/src/test/resources/jetty-logging.properties b/jetty-websocket/websocket-server/src/test/resources/jetty-logging.properties
index f64a66e3c3..4cadfe4a82 100644
--- a/jetty-websocket/websocket-server/src/test/resources/jetty-logging.properties
+++ b/jetty-websocket/websocket-server/src/test/resources/jetty-logging.properties
@@ -1,7 +1,9 @@
 org.eclipse.jetty.util.log.class=org.eclipse.jetty.util.log.StdErrLog
 org.eclipse.jetty.io.LEVEL=WARN
 org.eclipse.jetty.server.LEVEL=WARN
+org.eclipse.jetty.websocket.LEVEL=WARN
 org.eclipse.jetty.websocket.server.helper.RFCSocket.LEVEL=OFF
+# org.eclipse.jetty.websocket.io.LEVEL=DEBUG
 # org.eclipse.jetty.util.thread.QueuedThreadPool.LEVEL=DEBUG
 # org.eclipse.jetty.io.SelectorManager.LEVEL=INFO
 # org.eclipse.jetty.websocket.LEVEL=DEBUG