async def test_close_eofstream(loop: Any, ws_key: Any, key_data: Any) -> None:
resp = mock.Mock()
resp.status = 101
resp.headers = {
hdrs.UPGRADE: "websocket",
hdrs.CONNECTION: "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.request") as m_req:
m_os.urandom.return_value = key_data
m_req.return_value = loop.create_future()
m_req.return_value.set_result(resp)
writer = WebSocketWriter.return_value = mock.Mock()
session = aiohttp.ClientSession()
resp = await session.ws_connect("http://test.org")
assert not resp.closed
exc = EofStream()
resp._reader.set_exception(exc)
await resp.receive()
writer.close.assert_called_with(1000, b"")
assert resp.closed
await session.close()
async def test_close_eofstream(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 = loop.create_future()
m_req.return_value.set_result(resp)
writer = WebSocketWriter.return_value = mock.Mock()
session = aiohttp.ClientSession(loop=loop)
resp = await session.ws_connect('http://test.org')
assert not resp.closed
exc = EofStream()
resp._reader.set_exception(exc)
await resp.receive()
writer.close.assert_called_with(1000, b'')
assert resp.closed
await session.close()
async def test_receive_eofstream_in_reader(make_request, loop):
req = make_request('GET', '/')
ws = WebSocketResponse()
await ws.prepare(req)
ws._reader = mock.Mock()
exc = EofStream()
res = loop.create_future()
res.set_exception(exc)
ws._reader.read = make_mocked_coro(res)
ws._payload_writer.drain = mock.Mock()
ws._payload_writer.drain.return_value = loop.create_future()
ws._payload_writer.drain.return_value.set_result(True)
msg = await ws.receive()
assert msg.type == WSMsgType.CLOSED
assert ws.closed