1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167
|
import time
import unittest
from unittest import mock
import pytest # noqa: F401
from wsproto.events import AcceptConnection, CloseConnection, TextMessage, \
BytesMessage, Ping
import simple_websocket
class SimpleWebSocketClientTestCase(unittest.TestCase):
def get_client(self, mock_wsconn, url, events=[], subprotocols=None,
headers=None):
mock_wsconn().events.side_effect = \
[iter(ev) for ev in
[[AcceptConnection()]] + events + [[CloseConnection(1000)]]]
mock_wsconn().send = lambda x: str(x).encode('utf-8')
return simple_websocket.Client.connect(url, subprotocols=subprotocols,
headers=headers)
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_make_client(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1')
assert client.sock == mock_socket()
assert client.receive_bytes == 4096
assert client.input_buffer == []
assert client.event.__class__.__name__ == 'Event'
client.sock.send.assert_called_with(
b"Request(host='example.com', target='/ws?a=1', extensions=[], "
b"extra_headers=[], subprotocols=[])")
assert not client.is_server
assert client.host == 'example.com'
assert client.port == 80
assert client.path == '/ws?a=1'
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_make_client_subprotocol(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1',
subprotocols='foo')
assert client.subprotocols == ['foo']
client.sock.send.assert_called_with(
b"Request(host='example.com', target='/ws?a=1', extensions=[], "
b"extra_headers=[], subprotocols=['foo'])")
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_make_client_subprotocols(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1',
subprotocols=['foo', 'bar'])
assert client.subprotocols == ['foo', 'bar']
client.sock.send.assert_called_with(
b"Request(host='example.com', target='/ws?a=1', extensions=[], "
b"extra_headers=[], subprotocols=['foo', 'bar'])")
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_make_client_headers(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1',
headers={'Foo': 'Bar'})
client.sock.send.assert_called_with(
b"Request(host='example.com', target='/ws?a=1', extensions=[], "
b"extra_headers=[('Foo', 'Bar')], subprotocols=[])")
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_make_client_headers2(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws?a=1',
headers=[('Foo', 'Bar'), ('Foo', 'Baz')])
client.sock.send.assert_called_with(
b"Request(host='example.com', target='/ws?a=1', extensions=[], "
b"extra_headers=[('Foo', 'Bar'), ('Foo', 'Baz')], "
b"subprotocols=[])")
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_send(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws')
while client.connected:
time.sleep(0.01)
with pytest.raises(simple_websocket.ConnectionClosed):
client.send('hello')
client.connected = True
client.send('hello')
mock_socket().send.assert_called_with(
b"TextMessage(data='hello', frame_finished=True, "
b"message_finished=True)")
client.connected = True
client.send(b'hello')
mock_socket().send.assert_called_with(
b"Message(data=b'hello', frame_finished=True, "
b"message_finished=True)")
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_receive(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[
[TextMessage('hello')],
[BytesMessage(b'hello')],
])
while client.connected:
time.sleep(0.01)
client.connected = True
assert client.receive() == 'hello'
assert client.receive() == b'hello'
assert client.receive(timeout=0) is None
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_receive_after_close(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[
[TextMessage('hello')],
])
while client.connected:
time.sleep(0.01)
assert client.receive() == 'hello'
with pytest.raises(simple_websocket.ConnectionClosed):
client.receive()
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_receive_ping(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[
[Ping(b'hello')],
])
while client.connected:
time.sleep(0.01)
mock_socket().send.assert_any_call(b"Pong(payload=b'hello')")
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_receive_empty(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.side_effect = [b'x', b'x', b'']
client = self.get_client(mock_wsconn, 'ws://example.com/ws', events=[
[TextMessage('hello')],
])
while client.connected:
time.sleep(0.01)
client.connected = True
assert client.receive() == 'hello'
assert client.receive(timeout=0) is None
@mock.patch('simple_websocket.ws.socket.socket')
@mock.patch('simple_websocket.ws.WSConnection')
def test_close(self, mock_wsconn, mock_socket):
mock_socket.return_value.recv.return_value = b'x'
client = self.get_client(mock_wsconn, 'ws://example.com/ws')
while client.connected:
time.sleep(0.01)
with pytest.raises(simple_websocket.ConnectionClosed):
client.close()
client.connected = True
client.close()
assert not client.connected
mock_socket().send.assert_called_with(
b'CloseConnection(code=<CloseReason.NORMAL_CLOSURE: 1000>, '
b'reason=None)')
|