Tests: improved handshake for websocket tests.

test/test_java_websockets.py

@@ -8,7 +8,7 @@ from unit.applications.websockets import TestApplicationWebsocket
 class TestJavaWebsockets(TestApplicationJava):
     prerequisites = {'modules': ['java']}
 
-    ws = TestApplicationWebsocket(True)
+    ws = TestApplicationWebsocket()
 
     def setUp(self):
         super().setUp()

test/unit/applications/websockets.py

@@ -1,3 +1,4 @@
+import re
 import random
 import base64
 import struct
@@ -32,11 +33,7 @@ class TestApplicationWebsocket(TestApplicationProto):
 
     def upgrade(self):
         key = self.key()
-
+        _, sock = self.get(
-        if self.preinit:
-            self.get()
-
-        resp, sock = self.get(
             headers={
                 'Host': 'localhost',
                 'Upgrade': 'websocket',
@@ -45,10 +42,21 @@ class TestApplicationWebsocket(TestApplicationProto):
                 'Sec-WebSocket-Protocol': 'chat',
                 'Sec-WebSocket-Version': 13,
             },
-            read_timeout=1,
+            no_recv=True,
             start=True,
         )
 
+        resp = ''
+        while select.select([sock], [], [], 30)[0]:
+            resp += sock.recv(4096).decode()
+
+            if (
+                re.search('101 Switching Protocols', resp)
+                and resp[-4:] == '\r\n\r\n'
+            ):
+                resp = self._resp_to_dict(resp)
+                break
+
         return (resp, sock, key)
 
     def apply_mask(self, data, mask):