def test_parse_with_mask(self) -> None:
raw = b'\x81\x85\xc6\ti\x8d\xael\x05\xe1\xa9'
frame = WebsocketFrame()
frame.parse(raw)
self.assertEqual(frame.fin, True)
self.assertEqual(frame.rsv1, False)
self.assertEqual(frame.rsv2, False)
self.assertEqual(frame.rsv3, False)
self.assertEqual(frame.opcode, 0x1)
self.assertEqual(frame.masked, True)
assert frame.mask is not None
self.assertEqual(frame.mask, b'\xc6\ti\x8d')
self.assertEqual(frame.payload_length, 5)
self.assertEqual(frame.data, b'hello')
def test_build_with_mask(self) -> None:
raw = b'\x81\x85\xc6\ti\x8d\xael\x05\xe1\xa9'
frame = WebsocketFrame()
frame.fin = True
frame.opcode = websocketOpcodes.TEXT_FRAME
frame.masked = True
frame.mask = b'\xc6\ti\x8d'
frame.data = b'hello'
self.assertEqual(frame.build(), raw)
def test_handshake(self, mock_connect: mock.Mock,
mock_b64encode: mock.Mock) -> None:
key = b'MySecretKey'
mock_b64encode.return_value = key
mock_connect.return_value.recv.return_value = \
build_websocket_handshake_response(
WebsocketFrame.key_to_accept(key))
_ = WebsocketClient(DEFAULT_IPV4_HOSTNAME, DEFAULT_PORT)
mock_connect.return_value.send.assert_called_with(
build_websocket_handshake_request(key))
def test_send_recv_frames_success(
self,
mock_connect: mock.Mock,
mock_selector: mock.Mock,
mock_b64encode: mock.Mock,
) -> None:
key = b'MySecretKey'
mock_b64encode.return_value = key
mock_connect.return_value.recv.side_effect = [
build_websocket_handshake_response(
WebsocketFrame.key_to_accept(key), ),
WebsocketFrame.text(b'world'),
]
def on_message(frame: WebsocketFrame) -> None:
assert frame.build() == WebsocketFrame.text(b'world')
client = WebsocketClient(
b'localhost',
DEFAULT_PORT,
on_message=on_message,
)
mock_selector.assert_called_once()
client.handshake()
client.queue(memoryview(WebsocketFrame.text(b'hello')))
mock_connect.return_value.send.assert_called_once()
mock_selector.return_value.select.side_effect = [
[
(mock.Mock(), selectors.EVENT_WRITE),
],
]
client.run_once()
self.assertEqual(mock_connect.return_value.send.call_count, 2)
mock_selector.return_value.select.side_effect = [
[
(mock.Mock(), selectors.EVENT_READ),
],
]
client.run_once()
def test_handshake(self, mock_connect: mock.Mock,
mock_b64encode: mock.Mock,
mock_gethostbyname: mock.Mock) -> None:
key = b'MySecretKey'
mock_b64encode.return_value = key
mock_gethostbyname.return_value = '127.0.0.1'
mock_connect.return_value.recv.return_value = \
build_websocket_handshake_response(
WebsocketFrame.key_to_accept(key))
client = WebsocketClient(b'localhost', DEFAULT_PORT)
mock_connect.return_value.send.assert_not_called()
client.handshake()
mock_connect.return_value.send.assert_called_with(
build_websocket_handshake_request(key))
:copyright: (c) 2013-present by Abhinav Singh and contributors.
:license: BSD, see LICENSE for more details.
"""
import time
import logging
from proxy.http.websocket import (
WebsocketFrame,
WebsocketClient,
websocketOpcodes,
)
# globals
client: WebsocketClient
last_dispatch_time: float
static_frame = memoryview(WebsocketFrame.text(b'hello'))
num_echos = 10
logger = logging.getLogger(__name__)
def on_message(frame: WebsocketFrame) -> None:
"""WebsocketClient on_message callback."""
global client, num_echos, last_dispatch_time
logger.info(
'Received %r after %d millisec' %
(frame.data, (time.time() - last_dispatch_time) * 1000), )
assert (frame.data == b'hello'
and frame.opcode == websocketOpcodes.TEXT_FRAME)
if num_echos > 0:
client.queue(static_frame)
def on_message(frame: WebsocketFrame) -> None:
assert frame.build() == WebsocketFrame.text(b'world')