def client(loop, url, name):
ws = yield from aiohttp.ws_connect(url + '/getCaseCount')
num_tests = int((yield from ws.receive()).data)
print('running %d cases' % num_tests)
yield from ws.close()
for i in range(1, num_tests + 1):
print('running test case:', i)
text_url = url + '/runCase?case=%d&agent=%s' % (i, name)
ws = yield from aiohttp.ws_connect(text_url)
while True:
msg = yield from ws.receive()
if msg.type == aiohttp.MsgType.text:
ws.send_str(msg.data)
elif msg.type == aiohttp.MsgType.binary:
ws.send_bytes(msg.data)
elif msg.type == aiohttp.MsgType.close:
yield from ws.close()
break
else:
break
url = url + '/updateReports?agent=%s' % name
ws = yield from aiohttp.ws_connect(url)
yield from ws.close()
def test_close(self, m_req, m_os, WebSocketWriter):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
}
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(resp)
writer = WebSocketWriter.return_value = mock.Mock()
reader = resp.connection.reader.set_parser.return_value = mock.Mock()
resp = self.loop.run_until_complete(
aiohttp.ws_connect('http://test.org', loop=self.loop))
self.assertFalse(resp.closed)
msg = websocket.Message(websocket.MSG_CLOSE, b'', b'')
reader.read.return_value = helpers.create_future(self.loop)
reader.read.return_value.set_result(msg)
res = self.loop.run_until_complete(resp.close())
writer.close.assert_called_with(1000, b'')
self.assertTrue(resp.closed)
self.assertTrue(res)
self.assertIsNone(resp.exception())
# idempotent
res = self.loop.run_until_complete(resp.close())
self.assertFalse(res)
self.assertEqual(writer.close.call_count, 1)
def _run(self):
# print(self.client_id, 'Connecting')
self.stats['handshake_started_at'] = event_loop.time()
ws = yield from aiohttp.ws_connect(self.url)
self.stats['handshake_done_at'] = event_loop.time() # asyncio.get_event_loop().time()
# print(self.client_id, 'Connected')
# print(self.client_id, 'Sending the message')
self.stats['message_sent_at'] = event_loop.time()
ws.send_str(self.make_new_message())
while not ws.closed:
msg = yield from ws.receive()
if msg.tp == aiohttp.MsgType.text:
if msg.data in self.messages:
self.stats['echo_message_received_at'] = event_loop.time()
# print(self.client_id, 'Received the echo')
yield from ws.close()
break
elif msg.tp == aiohttp.MsgType.closed:
break
elif msg.tp == aiohttp.MsgType.error:
break
def test_close_exc2(self, m_req, m_os, WebSocketWriter):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
}
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(resp)
writer = WebSocketWriter.return_value = mock.Mock()
resp.connection.reader.set_parser.return_value = mock.Mock()
resp = self.loop.run_until_complete(
aiohttp.ws_connect(
'http://test.org', loop=self.loop))
self.assertFalse(resp.closed)
exc = ValueError()
writer.close.side_effect = exc
self.loop.run_until_complete(resp.close())
self.assertTrue(resp.closed)
self.assertIs(resp.exception(), exc)
resp._closed = False
writer.close.side_effect = asyncio.CancelledError()
self.assertRaises(asyncio.CancelledError,
self.loop.run_until_complete, resp.close())
def test_reader_read_exception(self, m_req, m_os, WebSocketWriter):
hresp = mock.Mock()
hresp.status = 101
hresp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
}
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(hresp)
WebSocketWriter.return_value = mock.Mock()
reader = hresp.connection.reader.set_parser.return_value = mock.Mock()
resp = self.loop.run_until_complete(
aiohttp.ws_connect(
'http://test.org', loop=self.loop))
exc = ValueError()
reader.read.return_value = helpers.create_future(self.loop)
reader.read.return_value.set_exception(exc)
msg = self.loop.run_until_complete(resp.receive())
self.assertEqual(msg.tp, aiohttp.MsgType.error)
self.assertIs(resp.exception(), exc)
def start_client(loop, url, filename):
# send request
ws = yield from aiohttp.ws_connect(url, autoclose=False, autoping=False)
# handshake
ws.send_str(json.dumps({'type': 'handshake', 'client': 'sample', 'protocol' : 'nlu'}))
msg = yield from ws.receive()
print(msg.data)
# msg = yield from ws.receive()
# print(msg.data)
#for i in range(2):
ws.send_str(json.dumps({'type': 'nlu', 'command': 'start'}))
frame_size = 640
with open(filename, 'rb') as file:
while file.readable():
data = file.read(frame_size)
if len(data) < frame_size:
break
ws.send_bytes(data)
ws.send_str(json.dumps({'type': 'nlu', 'command': 'stop'}))
output = yield from ws.receive() #asyncio.sleep(5)
print(json.loads(output.data))
def _connect(self, conn_type, session, force_close, force_release, pool):
future = self._future_class()
ws = aiohttp.ws_connect(self._url,
connector=self._connector,
loop=self._loop)
if self._timeout:
future_conn = asyncio.wait_for(ws, self._timeout, loop=self._loop)
else:
future_conn = asyncio. async (ws, loop=self._loop)
def on_connect(f):
try:
conn = f.result()
# Need to figure out some errors
except Exception as e:
future.set_exception(e)
else:
resp = Response(conn, self._future_class, loop=self._loop)
gc = conn_type(resp, self._future_class, self._timeout,
self._username, self._password, self._loop,
force_close, pool, force_release, session)
future.set_result(gc)
future_conn.add_done_callback(on_connect)
return future
def test_close_exc2(loop, ws_key, key_data):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter:
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
writer = WebSocketWriter.return_value = mock.Mock()
resp.connection.reader.set_parser.return_value = mock.Mock()
resp = yield from aiohttp.ws_connect('http://test.org',
loop=loop)
assert not resp.closed
exc = ValueError()
writer.close.side_effect = exc
yield from resp.close()
assert resp.closed
assert resp.exception() is exc
resp._closed = False
writer.close.side_effect = asyncio.CancelledError()
with pytest.raises(asyncio.CancelledError):
yield from resp.close()
async def join_rtm(self, filters=None):
"""Join the real-time messaging service.
Arguments:
filters (:py:class:`dict`, optional): Dictionary mapping
message filters to the functions they should dispatch to.
Use a :py:class:`collections.OrderedDict` if precedence is
important; only one filter, the first match, will be
applied to each message.
"""
if filters is None:
filters = self.MESSAGE_FILTERS
url = await self.get_socket_url()
logger.debug('Connecting to %r', url)
async with ws_connect(url) as socket:
first_msg = await socket.receive()
self._validate_first_message(first_msg)
async for message in socket:
if message.tp == MsgType.text:
result = await self.handle_message(message, filters)
if result is not None:
logger.info(
'Sending message: %r',
truncate(result, max_len=50),
)
socket.send_str(result)
elif message.tp in (MsgType.closed, MsgType.error):
if not socket.closed:
await socket.close()
break
logger.info('Left real-time messaging.')
def test_reader_read_exception(ws_key, key_data, loop):
hresp = mock.Mock()
hresp.status = 101
hresp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter:
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(hresp)
WebSocketWriter.return_value = mock.Mock()
reader = mock.Mock()
hresp.connection.reader.set_parser.return_value = reader
resp = yield from aiohttp.ws_connect('http://test.org',
loop=loop)
exc = ValueError()
reader.read.return_value = helpers.create_future(loop)
reader.read.return_value.set_exception(exc)
msg = yield from resp.receive()
assert msg.type == aiohttp.MsgType.ERROR
assert msg.type is msg.tp
assert resp.exception() is exc
def test_send_data_after_close(ws_key, key_data, loop):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter:
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
WebSocketWriter.return_value = mock.Mock()
resp = yield from aiohttp.ws_connect('http://test.org',
loop=loop)
resp._closed = True
pytest.raises(RuntimeError, resp.ping)
pytest.raises(RuntimeError, resp.pong)
pytest.raises(RuntimeError, resp.send_str, 's')
pytest.raises(RuntimeError, resp.send_bytes, b'b')
pytest.raises(RuntimeError, resp.send_json, {})
def test_ws_connect_custom_response(loop, ws_key, key_data):
class CustomResponse(ClientWebSocketResponse):
def read(self, decode=False):
return 'customized!'
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
res = yield from aiohttp.ws_connect(
'http://test.org',
ws_response_class=CustomResponse,
loop=loop)
assert res.read() == 'customized!'
async def join_rtm(self, filters=None):
"""Join the real-time messaging service."""
if filters is None:
filters = self.MESSAGE_FILTERS
url = await self.get_socket_url()
log.debug('Connecting to %r', url)
async with ws_connect(url) as socket:
# Export the socket to the class:
self.socket = socket
first_msg = await socket.receive()
self._validate_first_message(first_msg)
async for message in socket:
if message.tp == MsgType.text:
result = await self.handle_message(message, filters)
if result is not None:
log.info(
'Sending message: %r',
truncate(result, max_len=50),
)
socket.send_str(result)
elif message.tp in (MsgType.closed, MsgType.error):
if not socket.closed:
await socket.close()
break
log.info('Left real-time messaging.')
def test_reader_read_exception(ws_key, key_data, loop):
hresp = mock.Mock()
hresp.status = 101
hresp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter:
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(hresp)
WebSocketWriter.return_value = mock.Mock()
reader = mock.Mock()
hresp.connection.reader.set_parser.return_value = reader
resp = yield from aiohttp.ws_connect('http://test.org',
loop=loop)
exc = ValueError()
reader.read.return_value = helpers.create_future(loop)
reader.read.return_value.set_exception(exc)
msg = yield from resp.receive()
assert msg.type == aiohttp.MsgType.ERROR
assert msg.type is msg.tp
assert resp.exception() is exc
async def join_rtm(self, filters=None):
"""Join the real-time messaging service.
Arguments:
filters (:py:class:`dict`, optional): Dictionary mapping
message filters to the functions they should dispatch to.
Use a :py:class:`collections.OrderedDict` if precedence is
important; only one filter, the first match, will be
applied to each message.
"""
if filters is None:
filters = self.MESSAGE_FILTERS
url = await self.get_socket_url()
logger.debug('Connecting to %r', url)
async with ws_connect(url) as socket:
first_msg = await socket.receive()
self._validate_first_message(first_msg)
async for message in socket:
if message.tp == MsgType.text:
result = await self.handle_message(message, filters)
if result is not None:
logger.info(
'Sending message: %r',
truncate(result, max_len=50),
)
socket.send_str(result)
elif message.tp in (MsgType.closed, MsgType.error):
if not socket.closed:
await socket.close()
break
logger.info('Left real-time messaging.')
def test_close_exc2(loop, ws_key, key_data):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter:
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
writer = WebSocketWriter.return_value = mock.Mock()
resp.connection.reader.set_parser.return_value = mock.Mock()
resp = yield from aiohttp.ws_connect('http://test.org',
loop=loop)
assert not resp.closed
exc = ValueError()
writer.close.side_effect = exc
yield from resp.close()
assert resp.closed
assert resp.exception() is exc
resp._closed = False
writer.close.side_effect = asyncio.CancelledError()
with pytest.raises(asyncio.CancelledError):
yield from resp.close()
def test_send_data_after_close(ws_key, key_data, loop):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter:
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
WebSocketWriter.return_value = mock.Mock()
resp = yield from aiohttp.ws_connect('http://test.org',
loop=loop)
resp._closed = True
pytest.raises(RuntimeError, resp.ping)
pytest.raises(RuntimeError, resp.pong)
pytest.raises(RuntimeError, resp.send_str, 's')
pytest.raises(RuntimeError, resp.send_bytes, b'b')
pytest.raises(RuntimeError, resp.send_json, {})
def test_close_exc(self, m_req, m_os, WebSocketWriter):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
}
m_os.urandom.return_value = self.key_data
m_req.return_value = asyncio.Future(loop=self.loop)
m_req.return_value.set_result(resp)
WebSocketWriter.return_value = mock.Mock()
reader = resp.connection.reader.set_parser.return_value = mock.Mock()
resp = self.loop.run_until_complete(
aiohttp.ws_connect(
'http://test.org', loop=self.loop))
self.assertFalse(resp.closed)
exc = ValueError()
reader.read.return_value = asyncio.Future(loop=self.loop)
reader.read.return_value.set_exception(exc)
self.loop.run_until_complete(resp.close())
self.assertTrue(resp.closed)
self.assertIs(resp.exception(), exc)
def test_live_data(self):
ws = yield from aiohttp.ws_connect('http://localhost:12345/livedata/power')
yield from self.connection.publish(CURRENT_COST_KEY, 'a message')
message = yield from asyncio.wait_for(ws.receive(), 2)
self.assertEqual('a message', message.data)
yield from ws.close()
def _connect(self,
conn_type,
session,
force_close,
force_release,
pool):
future = self._future_class()
ws = aiohttp.ws_connect(
self._url, connector=self._connector, loop=self._loop)
if self._timeout:
future_conn = asyncio.wait_for(ws, self._timeout, loop=self._loop)
else:
future_conn = asyncio.async(ws, loop=self._loop)
def on_connect(f):
try:
conn = f.result()
# Need to figure out some errors
except Exception as e:
future.set_exception(e)
else:
resp = Response(conn, self._future_class, loop=self._loop)
gc = conn_type(resp, self._future_class, self._timeout,
self._username, self._password, self._loop,
force_close, pool, force_release, session)
future.set_result(gc)
future_conn.add_done_callback(on_connect)
return future
def test_send_data_after_close(ws_key, key_data, loop, mocker):
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: ws_key,
}
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
resp = yield from aiohttp.ws_connect('http://test.org',
loop=loop)
resp._writer._closing = True
mocker.spy(ws_logger, 'warning')
for meth, args in ((resp.ping, ()),
(resp.pong, ()),
(resp.send_str, ('s',)),
(resp.send_bytes, (b'b',)),
(resp.send_json, ({},))):
meth(*args)
assert ws_logger.warning.called
ws_logger.warning.reset_mock()
def test_live_data(self):
ws = yield from aiohttp.ws_connect(
'http://localhost:12345/livedata/power')
yield from self.connection.publish(CURRENT_COST_KEY, 'a message')
message = yield from asyncio.wait_for(ws.receive(), 2)
self.assertEqual('a message', message.data)
yield from ws.close()
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url, loop=self.loop)
resp.send_str('ask')
msg = yield from resp.receive()
self.assertEqual(msg.data, 'ask/answer')
yield from resp.close()
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url, loop=self.loop)
resp.send_str('ask')
msg = yield from resp.receive()
self.assertEqual(msg.data, 'ask/answer')
yield from resp.close()
def go(loop):
srv = yield from start_server(loop)
with pytest.raises(WSServerHandshakeError):
client = yield from ws_connect("http://localhost:9000?token=ooo")
yield from client.close()
srv.close()
yield from srv.wait_closed()
def go(loop):
srv = yield from start_server(loop)
with pytest.raises(WSServerHandshakeError):
client = yield from ws_connect("http://localhost:9000?token=ooo")
yield from client.close()
srv.close()
yield from srv.wait_closed()
def go(loop):
srv = yield from start_server(loop)
client = yield from ws_connect("http://localhost:9000")
message = yield from client.receive()
assert message.data == "hello, world!"
yield from client.close()
srv.close()
yield from srv.wait_closed()
def go(loop):
srv = yield from start_server(loop)
client = yield from ws_connect("http://localhost:9000")
message = yield from client.receive()
assert message.data == "hello, world!"
yield from client.close()
srv.close()
yield from srv.wait_closed()
def _websocket(self, url, **params):
if not params:
params = {
'stdout': 1,
'stderr': 1,
'stream': 1
}
url = self._endpoint(url) + "?" + urllib.parse.urlencode(params)
ws = yield from aiohttp.ws_connect(url, connector=self.connector)
return ws
def go(loop):
srv = yield from start_server(loop)
client = yield from ws_connect("http://localhost:9000")
yield from pub()
message = yield from client.receive()
assert message.data == subscription.channel.name
yield from client.close()
srv.close()
yield from srv.wait_closed()
def go(loop):
srv = yield from start_server(loop)
client = yield from ws_connect("http://localhost:9000")
yield from pub()
message = yield from client.receive()
assert message.data == subscription.channel.name
yield from client.close()
srv.close()
yield from srv.wait_closed()
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url, loop=self.loop)
resp.send_bytes(b'ask')
closed = yield from resp.close()
self.assertTrue(closed)
self.assertTrue(resp.closed)
self.assertEqual(resp.close_code, 1000)
msg = yield from resp.receive()
self.assertEqual(msg.tp, aiohttp.MsgType.closed)
def go(loop):
srv = yield from start_server(loop)
uri = "http://localhost:9000?token=" + token
client = yield from ws_connect(uri)
yield from pub()
message_ = yield from client.receive()
assert message_.data == message.body
yield from client.close()
srv.close()
yield from srv.wait_closed()
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url, loop=self.loop)
resp.send_bytes(b'ask')
closed = yield from resp.close()
self.assertTrue(closed)
self.assertTrue(resp.closed)
self.assertEqual(resp.close_code, 1000)
msg = yield from resp.receive()
self.assertEqual(msg.tp, aiohttp.MsgType.closed)
def go(loop):
srv = yield from start_server(loop)
uri = "http://localhost:9000?token=" + token
client = yield from ws_connect(uri)
yield from pub()
message_ = yield from client.receive()
assert message_.data == message.body
yield from client.close()
srv.close()
yield from srv.wait_closed()
def test_ws_connect_with_origin(self, m_req, m_os):
resp = mock.Mock()
resp.status = 403
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(resp)
origin = "https://example.org/page.html"
with self.assertRaises(errors.WSServerHandshakeError):
self.loop.run_until_complete(aiohttp.ws_connect("http://test.org", loop=self.loop, origin=origin))
self.assertIn(hdrs.ORIGIN, m_req.call_args[1]["headers"])
self.assertEqual(m_req.call_args[1]["headers"][hdrs.ORIGIN], origin)
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(
url, timeout=0.2, autoclose=False, loop=self.loop)
resp.send_bytes(b'ask')
msg = yield from resp.receive()
self.assertEqual(msg.data, 'test')
self.assertEqual(msg.tp, aiohttp.MsgType.text)
msg = yield from resp.close()
self.assertTrue(resp.closed)
self.assertIsInstance(resp.exception(), asyncio.TimeoutError)
def test_ws_connect_err_conn(self, m_req, m_os):
resp = mock.Mock()
resp.status = 101
resp.headers = {hdrs.UPGRADE: hdrs.WEBSOCKET, hdrs.CONNECTION: "close", hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key}
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(resp)
with self.assertRaises(errors.WSServerHandshakeError) as ctx:
self.loop.run_until_complete(
aiohttp.ws_connect("http://test.org", protocols=("t1", "t2", "chat"), loop=self.loop)
)
self.assertEqual(ctx.exception.message, "Invalid connection header")
def function2505(arg2110, arg533, arg2385):
var1538 = yield from aiohttp.ws_connect((arg533 + '/getCaseCount'))
var2751 = int(yield from var1538.receive().data)
print(('running %d cases' % var2751))
yield from var1538.close()
for var3246 in range(1, (var2751 + 1)):
print('running test case:', var3246)
var3727 = (arg533 + ('/runCase?case=%d&agent=%s' % (var3246, arg2385)))
var1538 = yield from aiohttp.ws_connect(var3727)
while True:
var3651 = yield from var1538.receive()
if (var3651.type == aiohttp.WSMsgType.text):
yield from var1538.send_str(var3651.data)
elif (var3651.type == aiohttp.WSMsgType.binary):
yield from var1538.send_bytes(var3651.data)
elif (var3651.type == aiohttp.WSMsgType.close):
yield from var1538.close()
break
else:
break
arg533 = (arg533 + ('/updateReports?agent=%s' % arg2385))
var1538 = yield from aiohttp.ws_connect(arg533)
yield from var1538.close()
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url, loop=self.loop)
resp.ping()
resp.send_bytes(b'ask')
msg = yield from resp.receive()
self.assertEqual(msg.tp, aiohttp.MsgType.binary)
self.assertEqual(msg.data, b'ask/answer')
msg = yield from resp.receive()
self.assertEqual(msg.tp, aiohttp.MsgType.close)
yield from resp.close()
yield from closed
def websocket_query(self, path, params={}):
"""
Open a websocket connection
:param path: Endpoint in API
:param params: Parameters added as a query arg
:returns: Websocket
"""
url = "http://docker/v" + self._api_version + "/" + path
connection = yield from aiohttp.ws_connect(url,
connector=self.connector(),
origin="http://docker",
autoping=True)
return connection
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(
url, autoclose=False, loop=self.loop)
resp.send_bytes(b'ask')
text = yield from resp.receive()
self.assertEqual(text.data, 'test')
t = asyncio.async(resp.close(), loop=self.loop)
yield from asyncio.sleep(0.1, loop=self.loop)
t.cancel()
yield from asyncio.sleep(0.1, loop=self.loop)
self.assertTrue(resp.closed)
self.assertIsNone(resp.exception())
def websocket_query(self, path, params={}):
"""
Open a websocket connection
:param path: Endpoint in API
:param params: Parameters added as a query arg
:returns: Websocket
"""
url = "http://docker/" + path
connection = yield from aiohttp.ws_connect(url,
connector=(yield from self.connector()),
origin="http://docker",
autoping=True)
return connection
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url,
timeout=0.2,
autoclose=False,
loop=self.loop)
resp.send_bytes(b'ask')
msg = yield from resp.receive()
self.assertEqual(msg.data, 'test')
self.assertEqual(msg.tp, aiohttp.MsgType.text)
msg = yield from resp.close()
self.assertTrue(resp.closed)
self.assertIsInstance(resp.exception(), asyncio.TimeoutError)
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url, loop=self.loop)
resp.ping()
resp.send_bytes(b'ask')
msg = yield from resp.receive()
self.assertEqual(msg.tp, aiohttp.MsgType.binary)
self.assertEqual(msg.data, b'ask/answer')
msg = yield from resp.receive()
self.assertEqual(msg.tp, aiohttp.MsgType.close)
yield from resp.close()
yield from closed
def test_ws_connect_with_origin(self, m_req, m_os):
resp = mock.Mock()
resp.status = 403
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(resp)
origin = 'https://example.org/page.html'
with self.assertRaises(errors.WSServerHandshakeError):
self.loop.run_until_complete(
aiohttp.ws_connect('http://test.org',
loop=self.loop,
origin=origin))
self.assertIn(hdrs.ORIGIN, m_req.call_args[1]["headers"])
self.assertEqual(m_req.call_args[1]["headers"][hdrs.ORIGIN], origin)
def go():
_, _, url = yield from self.create_server('GET', '/', handler)
resp = yield from aiohttp.ws_connect(url,
autoclose=False,
loop=self.loop)
resp.send_bytes(b'ask')
text = yield from resp.receive()
self.assertEqual(text.data, 'test')
t = asyncio. async (resp.close(), loop=self.loop)
yield from asyncio.sleep(0.1, loop=self.loop)
t.cancel()
yield from asyncio.sleep(0.1, loop=self.loop)
self.assertTrue(resp.closed)
self.assertIsNone(resp.exception())
def test_ws_connect_with_origin(key_data, loop):
resp = mock.Mock()
resp.status = 403
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
origin = 'https://example.org/page.html'
with pytest.raises(errors.WSServerHandshakeError):
yield from aiohttp.ws_connect('http://test.org',
loop=loop,
origin=origin)
assert hdrs.ORIGIN in m_req.call_args[1]["headers"]
assert m_req.call_args[1]["headers"][hdrs.ORIGIN] == origin
def go(loop):
srv = yield from start_server(loop)
uri = "http://localhost:9000?token=" + token
client = yield from ws_connect(uri)
yield from asyncio.sleep(1)
message.save()
message_ = yield from client.receive()
data = json.loads(message_.data)
message.refresh_from_db()
assert data[0]["pk"] == message.pk
yield from client.close()
srv.close()
yield from srv.wait_closed()
def test_ws_connect_close_resp_on_err(self, m_req, m_os):
resp = mock.Mock()
resp.status = 500
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key,
}
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(resp)
with self.assertRaises(errors.WSServerHandshakeError):
self.loop.run_until_complete(
aiohttp.ws_connect("http://test.org", protocols=("t1", "t2", "chat"), loop=self.loop)
)
resp.close.assert_called_with()
def test_ws_connect_with_origin(key_data, loop):
resp = mock.Mock()
resp.status = 403
with mock.patch('aiohttp.client.os') as m_os:
with mock.patch('aiohttp.client.ClientSession.get') as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = helpers.create_future(loop)
m_req.return_value.set_result(resp)
origin = 'https://example.org/page.html'
with pytest.raises(errors.WSServerHandshakeError):
yield from aiohttp.ws_connect('http://test.org',
loop=loop,
origin=origin)
assert hdrs.ORIGIN in m_req.call_args[1]["headers"]
assert m_req.call_args[1]["headers"][hdrs.ORIGIN] == origin
def go(loop):
srv = yield from start_server(loop)
uri = "http://localhost:9000?token=" + token
client = yield from ws_connect(uri)
yield from asyncio.sleep(1)
message.save()
message_ = yield from client.receive()
data = json.loads(message_.data)
message.refresh_from_db()
assert data[0]["pk"] == message.pk
yield from client.close()
srv.close()
yield from srv.wait_closed()
def test_ws_connect_close_resp_on_err(self, m_req, m_os):
resp = mock.Mock()
resp.status = 500
resp.headers = {
hdrs.UPGRADE: hdrs.WEBSOCKET,
hdrs.CONNECTION: hdrs.UPGRADE,
hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key
}
m_os.urandom.return_value = self.key_data
m_req.return_value = helpers.create_future(self.loop)
m_req.return_value.set_result(resp)
with self.assertRaises(errors.WSServerHandshakeError):
self.loop.run_until_complete(
aiohttp.ws_connect('http://test.org',
protocols=('t1', 't2', 'chat'),
loop=self.loop))
resp.close.assert_called_with()
