diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/test_java_websockets.py | 2 | ||||
| -rw-r--r-- | test/unit/applications/websockets.py | 20 |
2 files changed, 15 insertions, 7 deletions
diff --git a/test/test_java_websockets.py b/test/test_java_websockets.py index 3f2c0a8a..33091705 100644 --- a/test/test_java_websockets.py +++ b/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() diff --git a/test/unit/applications/websockets.py b/test/unit/applications/websockets.py index 50ff2797..5f78498c 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 @@ -32,11 +33,7 @@ class TestApplicationWebsocket(TestApplicationProto): def upgrade(self): key = self.key() - - if self.preinit: - self.get() - - resp, sock = self.get( + _, 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): |