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()
: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')