def test_message_async_handler(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
mock_server.async_handlers = True
_run(s.receive(packet.Packet(packet.MESSAGE, data='foo')))
mock_server._trigger_event.mock.assert_called_once_with('message',
'sid',
'foo',
run_async=True)
def test_upgrade_not_supported(self):
mock_server = self._get_mock_server()
mock_server._async['websocket'] = None
mock_server._async['websocket_class'] = None
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
environ = "foo"
_run(s._upgrade_websocket(environ))
mock_server._bad_request.assert_called_once_with()
def test_upgrade_handshake(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'foo')
s._upgrade_websocket = AsyncMock()
environ = {'REQUEST_METHOD': 'GET', 'QUERY_STRING': 'sid=foo',
'HTTP_CONNECTION': 'Foo,Upgrade,Bar',
'HTTP_UPGRADE': 'websocket'}
_run(s.handle_get_request(environ))
s._upgrade_websocket.mock.assert_called_once_with(environ)
def test_timeout(self):
mock_server = self._get_mock_server()
mock_server.ping_interval = 6
mock_server.ping_interval_grace_period = 2
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.last_ping = time.time() - 9
s.close = AsyncMock()
_run(s.send('packet'))
s.close.mock.assert_called_once_with(wait=False, abort=False)
def test_close_disconnect_error(self):
mock_server = self._get_mock_server()
mock_server._trigger_event.mock.side_effect = ZeroDivisionError
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
self.assertRaises(ZeroDivisionError, _run, s.close(wait=False))
self.assertTrue(s.closed)
self.assertEqual(mock_server._trigger_event.mock.call_count, 1)
mock_server._trigger_event.mock.assert_called_once_with(
'disconnect', 'sid')
def test_upgrade_twice(self):
mock_server = self._get_mock_server()
mock_server._async['websocket'] = mock.MagicMock()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
s.upgraded = True
environ = "foo"
with pytest.raises(IOError):
_run(s._upgrade_websocket(environ))
def test_polling_read(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'foo')
pkt1 = packet.Packet(packet.MESSAGE, data='hello')
pkt2 = packet.Packet(packet.MESSAGE, data='bye')
_run(s.send(pkt1))
_run(s.send(pkt2))
environ = {'REQUEST_METHOD': 'GET', 'QUERY_STRING': 'sid=foo'}
packets = _run(s.handle_get_request(environ))
self.assertEqual(packets, [pkt1, pkt2])
def test_upgrade_no_probe(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
ws = mock.MagicMock()
ws.wait = AsyncMock()
ws.wait.mock.return_value = packet.Packet(
packet.NOOP).encode(always_bytes=False)
_run(s._websocket_handler(ws))
self.assertFalse(s.upgraded)
def test_close_after_close(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
_run(s.close(wait=False))
self.assertTrue(s.closed)
self.assertEqual(mock_server._trigger_event.mock.call_count, 1)
mock_server._trigger_event.mock.assert_called_once_with(
'disconnect', 'sid')
_run(s.close())
self.assertEqual(mock_server._trigger_event.mock.call_count, 1)
def test_create(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
assert s.server == mock_server
assert s.sid == 'sid'
assert not s.upgraded
assert not s.closed
assert hasattr(s.queue, 'get')
assert hasattr(s.queue, 'put')
assert hasattr(s.queue, 'task_done')
assert hasattr(s.queue, 'join')
def test_create(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
self.assertEqual(s.server, mock_server)
self.assertEqual(s.sid, 'sid')
self.assertFalse(s.upgraded)
self.assertFalse(s.closed)
self.assertTrue(hasattr(s.queue, 'get'))
self.assertTrue(hasattr(s.queue, 'put'))
self.assertTrue(hasattr(s.queue, 'task_done'))
self.assertTrue(hasattr(s.queue, 'join'))
def test_upgrade(self):
mock_server = self._get_mock_server()
mock_server._async['websocket'] = mock.MagicMock()
mock_ws = AsyncMock()
mock_server._async['websocket'].return_value = mock_ws
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
environ = "foo"
_run(s._upgrade_websocket(environ))
mock_server._async['websocket'].assert_called_once_with(
s._websocket_handler)
mock_ws.mock.assert_called_once_with(environ)
async def _handle_connect(self,
environ,
transport,
b64=False,
jsonp_index=None):
"""Handle a client connection request."""
if self.start_service_task:
# start the service task to monitor connected clients
self.start_service_task = False
self.start_background_task(self._service_task)
sid = self._generate_id()
s = asyncio_socket.AsyncSocket(self, sid)
s.receive = types.MethodType(receive, s)
s.on_terminal = False
self.sockets[sid] = s
pkt = packet.Packet(
packet.OPEN, {
'sid': sid,
'upgrades': self._upgrades(sid, transport),
'pingTimeout': int(self.ping_timeout * 1000),
'pingInterval': int(self.ping_interval * 1000)
})
await s.send(pkt)
ret = await self._trigger_event('connect', sid, environ, run_async=False)
if ret is not None and ret is not True:
del self.sockets[sid]
self.logger.warning('Application rejected connection')
return self._unauthorized(ret or None)
if transport == 'websocket':
ret = await s.handle_get_request(environ)
if s.closed:
# websocket connection ended, so we are done
del self.sockets[sid]
return ret
else:
s.connected = True
headers = None
if self.cookie:
headers = [('Set-Cookie', self.cookie + '=' + sid)]
try:
return self._ok(await s.poll(),
headers=headers,
b64=b64,
jsonp_index=jsonp_index)
except exceptions.QueueEmpty as e:
print("_handle_connect", e)
return self._bad_request()
def test_polling_write(self):
mock_server = self._get_mock_server()
mock_server.max_http_buffer_size = 1000
pkt1 = packet.Packet(packet.MESSAGE, data='hello')
pkt2 = packet.Packet(packet.MESSAGE, data='bye')
p = payload.Payload(packets=[pkt1, pkt2]).encode()
s = asyncio_socket.AsyncSocket(mock_server, 'foo')
s.receive = AsyncMock()
environ = {'REQUEST_METHOD': 'POST', 'QUERY_STRING': 'sid=foo',
'CONTENT_LENGTH': len(p),
'wsgi.input': self._get_read_mock_coro(p)}
_run(s.handle_post_request(environ))
self.assertEqual(s.receive.mock.call_count, 2)
def test_schedule_ping(self):
mock_server = self._get_mock_server()
mock_server.ping_interval = 0.01
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.send = AsyncMock()
async def schedule_ping():
s.schedule_ping()
await asyncio.sleep(0.05)
_run(schedule_ping())
assert s.last_ping is not None
assert s.send.mock.call_args_list[0][0][0].encode() == '2'
def test_polling_write_too_large(self):
mock_server = self._get_mock_server()
pkt1 = packet.Packet(packet.MESSAGE, data='hello')
pkt2 = packet.Packet(packet.MESSAGE, data='bye')
p = payload.Payload(packets=[pkt1, pkt2]).encode()
mock_server.max_http_buffer_size = len(p) - 1
s = asyncio_socket.AsyncSocket(mock_server, 'foo')
s.receive = AsyncMock()
environ = {'REQUEST_METHOD': 'POST', 'QUERY_STRING': 'sid=foo',
'CONTENT_LENGTH': len(p),
'wsgi.input': self._get_read_mock_coro(p)}
self.assertRaises(exceptions.ContentTooLongError, _run,
s.handle_post_request(environ))
def test_schedule_ping_closed_socket(self):
mock_server = self._get_mock_server()
mock_server.ping_interval = 0.01
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.send = AsyncMock()
s.closed = True
async def schedule_ping():
s.schedule_ping()
await asyncio.sleep(0.05)
_run(schedule_ping())
assert s.last_ping is None
s.send.mock.assert_not_called()
def test_websocket_upgrade_with_payload(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
s.queue.join = AsyncMock(return_value=None)
probe = six.text_type('probe')
ws = mock.MagicMock()
ws.send = AsyncMock()
ws.wait = AsyncMock()
ws.wait.mock.side_effect = [
packet.Packet(packet.PING, data=probe).encode(always_bytes=False),
packet.Packet(packet.UPGRADE, data=b'2').encode(always_bytes=False)
]
_run(s._websocket_handler(ws))
self.assertTrue(s.upgraded)
def test_upgrade_no_upgrade_packet(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
s.queue.join = AsyncMock(return_value=None)
ws = mock.MagicMock()
ws.send = AsyncMock()
ws.wait = AsyncMock()
probe = 'probe'
ws.wait.mock.side_effect = [
packet.Packet(packet.PING, data=probe).encode(),
packet.Packet(packet.NOOP).encode(),
]
_run(s._websocket_handler(ws))
ws.send.mock.assert_called_once_with(
packet.Packet(packet.PONG, data=probe).encode())
assert _run(s.queue.get()).packet_type == packet.NOOP
assert not s.upgraded
def test_upgrade_no_upgrade_packet(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
s.queue.join = AsyncMock(return_value=None)
ws = mock.MagicMock()
ws.send = AsyncMock()
ws.wait = AsyncMock()
probe = six.text_type('probe')
ws.wait.mock.side_effect = [
packet.Packet(packet.PING, data=probe).encode(always_bytes=False),
packet.Packet(packet.NOOP).encode(always_bytes=False)
]
_run(s._websocket_handler(ws))
ws.send.mock.assert_called_once_with(
packet.Packet(packet.PONG, data=probe).encode(always_bytes=False))
self.assertEqual(_run(s.queue.get()).packet_type, packet.NOOP)
self.assertFalse(s.upgraded)
def test_websocket_read_write_wait_fail(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = False
s.queue.join = AsyncMock(return_value=None)
foo = six.text_type('foo')
bar = six.text_type('bar')
s.poll = AsyncMock(side_effect=[[
packet.Packet(packet.MESSAGE, data=bar)
], [packet.Packet(packet.MESSAGE, data=bar)], exceptions.QueueEmpty])
ws = mock.MagicMock()
ws.send = AsyncMock()
ws.wait = AsyncMock()
ws.wait.mock.side_effect = [
packet.Packet(packet.MESSAGE, data=foo).encode(always_bytes=False),
RuntimeError
]
ws.send.mock.side_effect = [None, RuntimeError]
_run(s._websocket_handler(ws))
self.assertEqual(s.closed, True)
def test_websocket_upgrade_with_backlog(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
s.queue.join = AsyncMock(return_value=None)
probe = six.text_type('probe')
foo = six.text_type('foo')
ws = mock.MagicMock()
ws.send = AsyncMock()
ws.wait = AsyncMock()
ws.wait.mock.side_effect = [
packet.Packet(packet.PING, data=probe).encode(always_bytes=False),
packet.Packet(packet.UPGRADE, data=b'2').encode(always_bytes=False)
]
s.upgrading = True
_run(s.send(packet.Packet(packet.MESSAGE, data=foo)))
_run(s._websocket_handler(ws))
self.assertTrue(s.upgraded)
self.assertFalse(s.upgrading)
self.assertEqual(s.packet_backlog, [])
ws.send.mock.assert_called_with('4foo')
def test_upgrade_no_upgrade_packet(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = True
s.queue.join = AsyncMock(return_value=None)
ws = mock.MagicMock()
ws.send = AsyncMock()
ws.wait = AsyncMock()
probe = six.text_type('probe')
ws.wait.mock.side_effect = [
packet.Packet(packet.PING, data=probe).encode(
always_bytes=False),
packet.Packet(packet.NOOP).encode(always_bytes=False)]
_run(s._websocket_handler(ws))
ws.send.mock.assert_called_once_with(packet.Packet(
packet.PONG, data=probe).encode(always_bytes=False))
self.assertEqual(_run(s.queue.get()).packet_type, packet.NOOP)
self.assertFalse(s.upgraded)
self.assertTrue(s.poll_ended)
environ = {'REQUEST_METHOD': 'GET', 'QUERY_STRING': 'sid=sid'}
packets = _run(s.handle_get_request(environ))
self.assertEqual(len(packets), 1)
self.assertEqual(packets[0].packet_type, packet.NOOP)
def test_websocket_ignore_invalid_packet(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
s.connected = False
s.queue.join = AsyncMock(return_value=None)
foo = six.text_type('foo')
bar = six.text_type('bar')
s.poll = AsyncMock(side_effect=[
[packet.Packet(packet.MESSAGE, data=bar)], exceptions.QueueEmpty])
ws = mock.MagicMock()
ws.send = AsyncMock()
ws.wait = AsyncMock()
ws.wait.mock.side_effect = [
packet.Packet(packet.OPEN).encode(always_bytes=False),
packet.Packet(packet.MESSAGE, data=foo).encode(
always_bytes=False),
None]
_run(s._websocket_handler(ws))
self.assertTrue(s.connected)
self.assertEqual(mock_server._trigger_event.mock.call_count, 2)
mock_server._trigger_event.mock.assert_has_calls([
mock.call('message', 'sid', foo, run_async=False),
mock.call('disconnect', 'sid')])
ws.send.mock.assert_called_with('4bar')
def test_poll_none(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
_run(s.queue.put(None))
self.assertEqual(_run(s.poll()), [])
def test_empty_poll(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
self.assertRaises(exceptions.QueueEmpty, _run, s.poll())
def test_invalid_packet(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
self.assertRaises(exceptions.UnknownPacketError, _run,
s.receive(packet.Packet(packet.OPEN)))
def test_polling_read_error(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'foo')
environ = {'REQUEST_METHOD': 'GET', 'QUERY_STRING': 'sid=foo'}
self.assertRaises(exceptions.QueueEmpty, _run,
s.handle_get_request(environ))
def test_send_after_close(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
_run(s.close(wait=False))
self.assertRaises(exceptions.SocketIsClosedError, _run,
s.send(packet.Packet(packet.NOOP)))
def test_empty_poll(self):
mock_server = self._get_mock_server()
s = asyncio_socket.AsyncSocket(mock_server, 'sid')
with pytest.raises(exceptions.QueueEmpty):
_run(s.poll())
