Merged with the default branch.1.13.0-1

1 files changed, 22 insertions, 9 deletions

diff --git a/test/unit/applications/websockets.py b/test/unit/applications/websockets.py
index 50ff2797..ef16f433 100644
--- a/test/unit/applications/websockets.py
+++ b/test/unit/applications/websockets.py
@@ -1,3 +1,4 @@
+import re
 import random
 import base64
 import struct
@@ -30,25 +31,37 @@ class TestApplicationWebsocket(TestApplicationProto):
         sha1 = hashlib.sha1((key + GUID).encode()).digest()
         return base64.b64encode(sha1).decode()
 
-    def upgrade(self):
-        key = self.key()
+    def upgrade(self, headers=None):
+        key = None
 
-        if self.preinit:
-            self.get()
-
-        resp, sock = self.get(
-            headers={
+        if headers is None:
+            key = self.key()
+            headers = {
                 'Host': 'localhost',
                 'Upgrade': 'websocket',
                 'Connection': 'Upgrade',
                 'Sec-WebSocket-Key': key,
                 'Sec-WebSocket-Protocol': 'chat',
                 'Sec-WebSocket-Version': 13,
-            },
-            read_timeout=1,
+            }
+
+        _, sock = self.get(
+            headers=headers,
+            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):