diff --git a/aiohttp/client.py b/aiohttp/client.py index dbc397279df..b1f33196d86 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -1,15 +1,21 @@ """HTTP Client for asyncio.""" import asyncio -import http.cookies -import urllib.parse -import warnings +import base64 +import hashlib +import os import sys import traceback +import warnings +import http.cookies +import urllib.parse import aiohttp from .client_reqrep import ClientRequest, ClientResponse +from .errors import WSServerHandshakeError from .multidict import MultiDictProxy, MultiDict, CIMultiDict +from .websocket import WS_KEY, WebSocketParser, WebSocketWriter +from .websocket_client import ClientWebSocketResponse from . import hdrs @@ -24,9 +30,10 @@ class ClientSession: _source_traceback = None _connector = None - def __init__(self, *, connector=None, loop=None, - request_class=ClientRequest, response_class=ClientResponse, - cookies=None, headers=None, auth=None): + def __init__(self, *, connector=None, loop=None, cookies=None, + headers=None, auth=None, request_class=ClientRequest, + response_class=ClientResponse, + ws_response_class=ClientWebSocketResponse): if loop is None: loop = asyncio.get_event_loop() self._loop = loop @@ -58,6 +65,7 @@ def __init__(self, *, connector=None, loop=None, self._request_class = request_class self._response_class = response_class + self._ws_response_class = ws_response_class if PY_34: def __del__(self): @@ -167,6 +175,70 @@ def request(self, method, url, *, return resp + @asyncio.coroutine + def ws_connect(self, url, *, + protocols=(), + timeout=10.0, + autoclose=True, + autoping=True): + """Initiate websocket connection.""" + + sec_key = base64.b64encode(os.urandom(16)) + + headers = { + hdrs.UPGRADE: hdrs.WEBSOCKET, + hdrs.CONNECTION: hdrs.UPGRADE, + hdrs.SEC_WEBSOCKET_VERSION: '13', + hdrs.SEC_WEBSOCKET_KEY: sec_key.decode(), + } + if protocols: + headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = ','.join(protocols) + + # send request + resp = yield from self.request('get', url, headers=headers, + read_until_eof=False) + + # check handshake + if resp.status != 101: + raise WSServerHandshakeError('Invalid response status') + + if resp.headers.get(hdrs.UPGRADE, '').lower() != 'websocket': + raise WSServerHandshakeError('Invalid upgrade header') + + if resp.headers.get(hdrs.CONNECTION, '').lower() != 'upgrade': + raise WSServerHandshakeError('Invalid connection header') + + # key calculation + key = resp.headers.get(hdrs.SEC_WEBSOCKET_ACCEPT, '') + match = base64.b64encode( + hashlib.sha1(sec_key + WS_KEY).digest()).decode() + if key != match: + raise WSServerHandshakeError('Invalid challenge response') + + # websocket protocol + protocol = None + if protocols and hdrs.SEC_WEBSOCKET_PROTOCOL in resp.headers: + resp_protocols = [ + proto.strip() for proto in + resp.headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(',')] + + for proto in resp_protocols: + if proto in protocols: + protocol = proto + break + + reader = resp.connection.reader.set_parser(WebSocketParser) + writer = WebSocketWriter(resp.connection.writer, use_mask=True) + + return self._ws_response_class(reader, + writer, + protocol, + resp, + timeout, + autoclose, + autoping, + self._loop) + def _update_cookies(self, cookies): """Update shared cookies.""" if isinstance(cookies, dict): diff --git a/aiohttp/websocket_client.py b/aiohttp/websocket_client.py index 706cc50a5fd..718f6da4df9 100644 --- a/aiohttp/websocket_client.py +++ b/aiohttp/websocket_client.py @@ -1,14 +1,10 @@ """WebSocket client for asyncio.""" import asyncio -import base64 -import hashlib -import os - -from aiohttp import client, hdrs -from .errors import WSServerHandshakeError -from .websocket import WS_KEY, Message -from .websocket import WebSocketParser, WebSocketWriter, WebSocketError + +import aiohttp +from .websocket import Message +from .websocket import WebSocketError from .websocket import MSG_BINARY, MSG_TEXT, MSG_CLOSE, MSG_PING, MSG_PONG __all__ = ('ws_connect', 'MsgType') @@ -33,67 +29,6 @@ class MsgType(IntEnum): closedMessage = Message(MsgType.closed, None, None) -@asyncio.coroutine -def ws_connect(url, protocols=(), timeout=10.0, connector=None, - response_class=None, autoclose=True, autoping=True, loop=None): - """Initiate websocket connection.""" - if loop is None: - loop = asyncio.get_event_loop() - - sec_key = base64.b64encode(os.urandom(16)) - - headers = { - hdrs.UPGRADE: hdrs.WEBSOCKET, - hdrs.CONNECTION: hdrs.UPGRADE, - hdrs.SEC_WEBSOCKET_VERSION: '13', - hdrs.SEC_WEBSOCKET_KEY: sec_key.decode(), - } - if protocols: - headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = ','.join(protocols) - - # send request - resp = yield from client.request( - 'get', url, headers=headers, - read_until_eof=False, - connector=connector, loop=loop) - - # check handshake - if resp.status != 101: - raise WSServerHandshakeError('Invalid response status') - - if resp.headers.get(hdrs.UPGRADE, '').lower() != 'websocket': - raise WSServerHandshakeError('Invalid upgrade header') - - if resp.headers.get(hdrs.CONNECTION, '').lower() != 'upgrade': - raise WSServerHandshakeError('Invalid connection header') - - # key calculation - key = resp.headers.get(hdrs.SEC_WEBSOCKET_ACCEPT, '') - match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()).decode() - if key != match: - raise WSServerHandshakeError('Invalid challenge response') - - # websocket protocol - protocol = None - if protocols and hdrs.SEC_WEBSOCKET_PROTOCOL in resp.headers: - resp_protocols = [proto.strip() for proto in - resp.headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(',')] - - for proto in resp_protocols: - if proto in protocols: - protocol = proto - break - - reader = resp.connection.reader.set_parser(WebSocketParser) - writer = WebSocketWriter(resp.connection.writer, use_mask=True) - - if response_class is None: - response_class = ClientWebSocketResponse - - return response_class( - reader, writer, protocol, resp, timeout, autoclose, autoping, loop) - - class ClientWebSocketResponse: def __init__(self, reader, writer, protocol, @@ -236,3 +171,30 @@ def receive(self): return msg finally: self._waiting = False + + +@asyncio.coroutine +def ws_connect(url, *, protocols=(), timeout=10.0, connector=None, + ws_response_class=ClientWebSocketResponse, autoclose=True, + autoping=True, loop=None): + + if loop is None: + asyncio.get_event_loop() + + if connector is None: + connector = aiohttp.TCPConnector(loop=loop, force_close=True) + + session = aiohttp.ClientSession(loop=loop, connector=connector, + ws_response_class=ws_response_class) + + try: + resp = yield from session.ws_connect( + url, + protocols=protocols, + timeout=timeout, + autoclose=autoclose, + autoping=autoping) + return resp + + finally: + session.detach() diff --git a/docs/client_reference.rst b/docs/client_reference.rst index 915dd7356b4..757302dc1f2 100644 --- a/docs/client_reference.rst +++ b/docs/client_reference.rst @@ -29,9 +29,9 @@ Usage example:: .. versionadded:: 0.15.2 -.. class:: ClientSession(*, connector=None, loop=None, request_class=None,\ - response_class=None, cookies=None, headers=None,\ - auth=None) +.. class:: ClientSession(*, connector=None, loop=None, cookies=None,\ + headers=None, auth=None, request_class=ClientRequest,\ + response_class=ClientResponse, ws_response_class=ClientWebSocketResponse) The class for creating client sessions and making requests. @@ -46,11 +46,6 @@ Usage example:: recommend to use explicit loops everywhere. (optional) - - :param request_class: Custom Request class implementation (optional) - - :param response_class: Custom Response class implementation (optional) - :param dict cookies: Cookies to send with the request (optional) :param dict headers: HTTP Headers to send with @@ -59,6 +54,23 @@ Usage example:: :param aiohttp.helpers.BasicAuth auth: BasicAuth named tuple that represents HTTP Basic Auth (optional) + :param request_class: Request class implementation. ``ClientRequest`` by + default. + + :param response_class: Response class implementation. ``ClientResponse`` by + default. + + :param ws_response_class: WebSocketResponse class implementation. + ``ClientWebSocketResponse`` by default. + + .. versionadded:: 0.16 + + .. versionchanged:: 0.16 + *request_class* default changed from ``None`` to ``ClientRequest`` + + .. versionchanged:: 0.16 + *response_class* default changed from ``None`` to ``ClientResponse`` + .. attribute:: closed ``True`` if the session has been closed, ``False`` otherwise. @@ -223,6 +235,26 @@ Usage example:: :param data: Dictionary, bytes, or file-like object to send in the body of the request (optional) + + .. coroutinemethod:: ws_connect(url, *, protocols=(), timeout=10.0\ + autoclose=True, autoping=True) + + Create a websocket connection. Returns a :class:`ClientWebSocketResponse` object. + + :param str url: Websocket server url + + :param tuple protocols: Websocket protocols + + :param float timeout: Timeout for websocket read. 10 seconds by default + + :param bool autoclose: Automatically close websocket connection on close + message from server. If `autoclose` is False + them close procedure has to be handled manually + + :param bool autoping: automatically send `pong` on `ping` message from server + + .. versionadded:: 0.16 + .. method:: close() Close underlying connector. diff --git a/docs/client_websockets.rst b/docs/client_websockets.rst index 9bd6c4e21ab..b771c6ecc46 100644 --- a/docs/client_websockets.rst +++ b/docs/client_websockets.rst @@ -46,25 +46,32 @@ ClientWebSocketResponse To connect to a websocket server you have to use the `aiohttp.ws_connect()` function, do not create an instance of class :class:`ClientWebSocketResponse` manually. -.. py:function:: ws_connect(url, protocols=(), connector=None, response_class=None, autoclose=True, autoping=True, loop=None) +.. coroutinefunction:: ws_connect(url, *, protocols=(), timeout=10.0, connector=None,\ + ws_response_class=ClientWebSocketResponse,\ + autoclose=True, autoping=True, loop=None) This function creates a websocket connection, checks the response and returns a :class:`ClientWebSocketResponse` object. In case of failure it may raise a :exc:`~aiohttp.errors.WSServerHandshakeError` exception. - :param str url: websocket server url + :param str url: Websocket server url - :param tuple protocols: websocket protocols + :param tuple protocols: Websocket protocols + + :param float timeout: Timeout for websocket read. 10 seconds by default :param obj connector: object :class:`TCPConnector` - :param response_class: (optional) Custom Response class implementation. + :param ws_response_class: WebSocketResponse class implementation. + ``ClientWebSocketResponse`` by default. + + .. versionadded:: 0.16 - :param bool autoclose: automatically close websocket connection - on close message from server. if `autoclose` is + :param bool autoclose: Automatically close websocket connection + on close message from server. If `autoclose` is False them close procedure has to be handled manually - :param bool autoping: automatically send `pong` on `ping` message from server + :param bool autoping: Automatically send `pong` on `ping` message from server :param loop: :ref:`event loop` used for processing HTTP requests. diff --git a/tests/test_websocket_client.py b/tests/test_websocket_client.py index 38ec9493cf8..debb58ca215 100644 --- a/tests/test_websocket_client.py +++ b/tests/test_websocket_client.py @@ -24,9 +24,9 @@ def setUp(self): def tearDown(self): self.loop.close() - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_ws_connect(self, m_client, m_os): + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_ws_connect(self, m_req, m_os): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -36,8 +36,8 @@ def test_ws_connect(self, m_client, m_os): hdrs.SEC_WEBSOCKET_PROTOCOL: 'chat' } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) res = self.loop.run_until_complete( websocket_client.ws_connect( @@ -48,9 +48,9 @@ def test_ws_connect(self, m_client, m_os): self.assertIsInstance(res, websocket_client.ClientWebSocketResponse) self.assertEqual(res.protocol, 'chat') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_ws_connect_custom_response(self, m_client, m_os): + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_ws_connect_custom_response(self, m_req, m_os): class CustomResponse(websocket_client.ClientWebSocketResponse): def read(self, decode=False): @@ -64,20 +64,20 @@ def read(self, decode=False): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key, } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) res = self.loop.run_until_complete( websocket_client.ws_connect( 'http://test.org', - response_class=CustomResponse, + ws_response_class=CustomResponse, loop=self.loop)) self.assertEqual(res.read(), 'customized!') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_ws_connect_global_loop(self, m_client, m_os): + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_ws_connect_global_loop(self, m_req, m_os): asyncio.set_event_loop(self.loop) resp = mock.Mock() @@ -88,8 +88,8 @@ def test_ws_connect_global_loop(self, m_client, m_os): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) resp = self.loop.run_until_complete( websocket_client.ws_connect('http://test.org')) @@ -97,9 +97,9 @@ def test_ws_connect_global_loop(self, m_client, m_os): asyncio.set_event_loop(None) - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_ws_connect_err_status(self, m_client, m_os): + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_ws_connect_err_status(self, m_req, m_os): resp = mock.Mock() resp.status = 500 resp.headers = { @@ -108,8 +108,8 @@ def test_ws_connect_err_status(self, m_client, m_os): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) with self.assertRaises(errors.WSServerHandshakeError) as ctx: self.loop.run_until_complete( @@ -120,9 +120,9 @@ def test_ws_connect_err_status(self, m_client, m_os): self.assertEqual( ctx.exception.message, 'Invalid response status') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_ws_connect_err_upgrade(self, m_client, m_os): + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_ws_connect_err_upgrade(self, m_req, m_os): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -131,8 +131,8 @@ def test_ws_connect_err_upgrade(self, m_client, m_os): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) with self.assertRaises(errors.WSServerHandshakeError) as ctx: self.loop.run_until_complete( @@ -143,9 +143,9 @@ def test_ws_connect_err_upgrade(self, m_client, m_os): self.assertEqual( ctx.exception.message, 'Invalid upgrade header') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_ws_connect_err_conn(self, m_client, m_os): + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_ws_connect_err_conn(self, m_req, m_os): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -154,8 +154,8 @@ def test_ws_connect_err_conn(self, m_client, m_os): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) with self.assertRaises(errors.WSServerHandshakeError) as ctx: self.loop.run_until_complete( @@ -166,9 +166,9 @@ def test_ws_connect_err_conn(self, m_client, m_os): self.assertEqual( ctx.exception.message, 'Invalid connection header') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_ws_connect_err_challenge(self, m_client, m_os): + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_ws_connect_err_challenge(self, m_req, m_os): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -177,8 +177,8 @@ def test_ws_connect_err_challenge(self, m_client, m_os): hdrs.SEC_WEBSOCKET_ACCEPT: 'asdfasdfasdfasdfasdfasdf' } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) with self.assertRaises(errors.WSServerHandshakeError) as ctx: self.loop.run_until_complete( @@ -189,10 +189,10 @@ def test_ws_connect_err_challenge(self, m_client, m_os): self.assertEqual( ctx.exception.message, 'Invalid challenge response') - @mock.patch('aiohttp.websocket_client.WebSocketWriter') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_close(self, m_client, m_os, WebSocketWriter): + @mock.patch('aiohttp.client.WebSocketWriter') + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_close(self, m_req, m_os, WebSocketWriter): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -201,8 +201,8 @@ def test_close(self, m_client, m_os, WebSocketWriter): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key, } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=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() @@ -225,10 +225,10 @@ def test_close(self, m_client, m_os, WebSocketWriter): self.assertFalse(res) self.assertEqual(writer.close.call_count, 1) - @mock.patch('aiohttp.websocket_client.WebSocketWriter') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_close_exc(self, m_client, m_os, WebSocketWriter): + @mock.patch('aiohttp.client.WebSocketWriter') + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_close_exc(self, m_req, m_os, WebSocketWriter): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -237,8 +237,8 @@ def test_close_exc(self, m_client, m_os, WebSocketWriter): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key, } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + 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() @@ -255,10 +255,10 @@ def test_close_exc(self, m_client, m_os, WebSocketWriter): self.assertTrue(resp.closed) self.assertIs(resp.exception(), exc) - @mock.patch('aiohttp.websocket_client.WebSocketWriter') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_close_exc2(self, m_client, m_os, WebSocketWriter): + @mock.patch('aiohttp.client.WebSocketWriter') + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_close_exc2(self, m_req, m_os, WebSocketWriter): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -267,8 +267,8 @@ def test_close_exc2(self, m_client, m_os, WebSocketWriter): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key, } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) writer = WebSocketWriter.return_value = mock.Mock() resp.connection.reader.set_parser.return_value = mock.Mock() @@ -289,10 +289,10 @@ def test_close_exc2(self, m_client, m_os, WebSocketWriter): self.assertRaises(asyncio.CancelledError, self.loop.run_until_complete, resp.close()) - @mock.patch('aiohttp.websocket_client.WebSocketWriter') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_send_data_after_close(self, m_client, m_os, WebSocketWriter): + @mock.patch('aiohttp.client.WebSocketWriter') + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_send_data_after_close(self, m_req, m_os, WebSocketWriter): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -301,8 +301,8 @@ def test_send_data_after_close(self, m_client, m_os, WebSocketWriter): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key, } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) WebSocketWriter.return_value = mock.Mock() resp = self.loop.run_until_complete( @@ -315,10 +315,10 @@ def test_send_data_after_close(self, m_client, m_os, WebSocketWriter): self.assertRaises(RuntimeError, resp.send_str, 's') self.assertRaises(RuntimeError, resp.send_bytes, b'b') - @mock.patch('aiohttp.websocket_client.WebSocketWriter') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_send_data_type_errors(self, m_client, m_os, WebSocketWriter): + @mock.patch('aiohttp.client.WebSocketWriter') + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_send_data_type_errors(self, m_req, m_os, WebSocketWriter): resp = mock.Mock() resp.status = 101 resp.headers = { @@ -327,8 +327,8 @@ def test_send_data_type_errors(self, m_client, m_os, WebSocketWriter): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key, } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(resp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(resp) WebSocketWriter.return_value = mock.Mock() resp = self.loop.run_until_complete( @@ -338,10 +338,10 @@ def test_send_data_type_errors(self, m_client, m_os, WebSocketWriter): self.assertRaises(TypeError, resp.send_str, b's') self.assertRaises(TypeError, resp.send_bytes, 'b') - @mock.patch('aiohttp.websocket_client.WebSocketWriter') - @mock.patch('aiohttp.websocket_client.os') - @mock.patch('aiohttp.websocket_client.client') - def test_reader_read_exception(self, m_client, m_os, WebSocketWriter): + @mock.patch('aiohttp.client.WebSocketWriter') + @mock.patch('aiohttp.client.os') + @mock.patch('aiohttp.client.ClientSession.request') + def test_reader_read_exception(self, m_req, m_os, WebSocketWriter): hresp = mock.Mock() hresp.status = 101 hresp.headers = { @@ -350,8 +350,8 @@ def test_reader_read_exception(self, m_client, m_os, WebSocketWriter): hdrs.SEC_WEBSOCKET_ACCEPT: self.ws_key, } m_os.urandom.return_value = self.key_data - m_client.request.return_value = asyncio.Future(loop=self.loop) - m_client.request.return_value.set_result(hresp) + m_req.return_value = asyncio.Future(loop=self.loop) + m_req.return_value.set_result(hresp) WebSocketWriter.return_value = mock.Mock() reader = hresp.connection.reader.set_parser.return_value = mock.Mock()