1 files changed, 56 insertions, 7 deletions
diff --git a/test/test_websockets.py b/test/test_websockets.py
index 0b2647ef..a5ebf3d1 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -1,29 +1,78 @@
+from netlib import tcp
 from netlib import test
 from netlib.websockets import implementations as impl
 from netlib.websockets import websockets as ws
 import os
+from nose.tools import raises
 
 class TestWebSockets(test.ServerTestBase):
     handler = impl.WebSocketsEchoHandler
 
+    def random_bytes(self, n = 100):
+        return os.urandom(n)
+
     def echo(self, msg):
         client = impl.WebSocketsClient(("127.0.0.1", self.port))
         client.connect()
         client.send_message(msg)
         response = client.read_next_message()
-        print "Assert response: " + response + " == msg: " + msg
         assert response == msg
 
     def test_simple_echo(self):
         self.echo("hello I'm the client")
 
     def test_frame_sizes(self):
-        small_string     =  os.urandom(100)   # length can fit in the the 7 bit payload length 
-        medium_string    =  os.urandom(50000) # 50kb, sligthly larger than can fit in a 7 bit int
-        large_string     =  os.urandom(150000) # 150kb, slightly larger than can fit in a 16 bit int
+        small_msg  = self.random_bytes(100)   # length can fit in the the 7 bit payload length 
+        medium_msg = self.random_bytes(50000) # 50kb, sligthly larger than can fit in a 7 bit int
+        large_msg  = self.random_bytes(150000) # 150kb, slightly larger than can fit in a 16 bit int
+
+        self.echo(small_msg)
+        self.echo(medium_msg)
+        self.echo(large_msg)
+
+    def test_default_builder(self):
+        """
+          default builder should always generate valid frames
+        """ 
+        msg = self.random_bytes()
+        client_frame = ws.WebSocketsFrame.default(msg, from_client = True)
+        assert client_frame.is_valid()
+
+        server_frame = ws.WebSocketsFrame.default(msg, from_client = False)
+        assert server_frame.is_valid()
+
+    def test_serialization_bijection(self):
+        for is_client in [True, False]:
+          for num_bytes in [100, 50000, 150000]:
+            frame = ws.WebSocketsFrame.default(self.random_bytes(num_bytes), is_client)
+            assert frame == ws.WebSocketsFrame.from_bytes(frame.to_bytes())
+
+        bytes = b'\x81\x11cba'
+        assert ws.WebSocketsFrame.from_bytes(bytes).to_bytes() == bytes
+
+
+class BadHandshakeHandler(impl.WebSocketsEchoHandler):
+    def handshake(self):
+        client_hs = ws.read_handshake(self.rfile.read, 1)
+        key       = ws.process_handshake_from_client(client_hs)
+        response  = ws.create_server_handshake("malformed_key")
+        self.wfile.write(response)
+        self.wfile.flush()
+        self.handshake_done = True
+
+class TestBadHandshake(test.ServerTestBase):
+    """
+      Ensure that the client disconnects if the server handshake is malformed
+    """ 
+    handler = BadHandshakeHandler
+
+    @raises(tcp.NetLibDisconnect)
+    def test(self):
+        client = impl.WebSocketsClient(("127.0.0.1", self.port))
+        client.connect()
+        client.send_message("hello")
+
+
 
-        self.echo(small_string)
-        self.echo(medium_string)
-        self.echo(large_string)