From 2545cef741946ad9b9c67228eff624cd432cc071 Mon Sep 17 00:00:00 2001 From: davebshow Date: Wed, 20 May 2015 08:54:41 -0400 Subject: [PATCH 1/5] cleaned up for PR --- aiohttp/client.py | 80 +++++++++++++++++- aiohttp/websocket_client.py | 93 +++++++-------------- docs/client_reference.rst | 22 +++++ docs/client_websockets.rst | 18 +++-- tests/test_websocket_client.py | 144 ++++++++++++++++----------------- 5 files changed, 209 insertions(+), 148 deletions(-) diff --git a/aiohttp/client.py b/aiohttp/client.py index dbc397279df..7b095a7d76c 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 @@ -167,6 +173,74 @@ def request(self, method, url, *, return resp + @asyncio.coroutine + def ws_connect(self, url, *, + protocols=(), + timeout=10.0, + ws_response_class=None, + 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) + + if ws_response_class is None: + ws_response_class = ClientWebSocketResponse + + return 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..8e3850ccc4c 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,65 +29,30 @@ 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.""" +def ws_connect(url, *, protocols=(), timeout=10.0, connector=None, + ws_response_class=None, autoclose=True, autoping=True, + loop=None): + 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) + asyncio.get_event_loop() + + if connector is None: + connector = aiohttp.TCPConnector(loop=loop, force_close=True) + + session = aiohttp.ClientSession(loop=loop, connector=connector) + + try: + resp = yield from session.ws_connect( + url, + protocols=protocols, + timeout=timeout, + ws_response_class=ws_response_class, + autoclose=autoclose, + autoping=autoping) + return resp + + finally: + session.detach() class ClientWebSocketResponse: diff --git a/docs/client_reference.rst b/docs/client_reference.rst index 915dd7356b4..0a225b79966 100644 --- a/docs/client_reference.rst +++ b/docs/client_reference.rst @@ -223,6 +223,28 @@ 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\ + ws_response_class=None, 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 ws_response_class: (optional) Custom Response class implementation + + :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 + + .. method:: close() Close underlying connector. diff --git a/docs/client_websockets.rst b/docs/client_websockets.rst index 9bd6c4e21ab..fc182c15058 100644 --- a/docs/client_websockets.rst +++ b/docs/client_websockets.rst @@ -46,25 +46,29 @@ 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) +.. py:function:: ws_connect(url, *, protocols=(), timeout=10.0,\ + connector=None, ws_response_class=None,\ + 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: (optional) Custom Response class implementation. - :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() From 29fee1d60809558fb37ed1d3942899796d5798fc Mon Sep 17 00:00:00 2001 From: davebshow Date: Thu, 21 May 2015 09:22:35 -0400 Subject: [PATCH 2/5] moved ws_response_class param to ClientSession constructor. updated docs --- aiohttp/client.py | 24 +++++++++++------------- aiohttp/websocket_client.py | 8 ++++++-- docs/client_reference.rst | 12 ++++++------ docs/client_websockets.rst | 3 ++- 4 files changed, 25 insertions(+), 22 deletions(-) diff --git a/aiohttp/client.py b/aiohttp/client.py index 7b095a7d76c..26b1eb46b4e 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -32,7 +32,8 @@ class ClientSession: def __init__(self, *, connector=None, loop=None, request_class=ClientRequest, response_class=ClientResponse, - cookies=None, headers=None, auth=None): + ws_response_class=ClientWebSocketResponse, cookies=None, + headers=None, auth=None): if loop is None: loop = asyncio.get_event_loop() self._loop = loop @@ -64,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): @@ -177,7 +179,6 @@ def request(self, method, url, *, def ws_connect(self, url, *, protocols=(), timeout=10.0, - ws_response_class=None, autoclose=True, autoping=True): """Initiate websocket connection.""" @@ -229,17 +230,14 @@ def ws_connect(self, url, *, reader = resp.connection.reader.set_parser(WebSocketParser) writer = WebSocketWriter(resp.connection.writer, use_mask=True) - if ws_response_class is None: - ws_response_class = ClientWebSocketResponse - - return ws_response_class(reader, - writer, - protocol, - resp, - timeout, - autoclose, - autoping, - self._loop) + return self._ws_response_class(reader, + writer, + protocol, + resp, + timeout, + autoclose, + autoping, + self._loop) def _update_cookies(self, cookies): """Update shared cookies.""" diff --git a/aiohttp/websocket_client.py b/aiohttp/websocket_client.py index 8e3850ccc4c..ea5841b367d 100644 --- a/aiohttp/websocket_client.py +++ b/aiohttp/websocket_client.py @@ -39,14 +39,18 @@ def ws_connect(url, *, protocols=(), timeout=10.0, connector=None, if connector is None: connector = aiohttp.TCPConnector(loop=loop, force_close=True) - session = aiohttp.ClientSession(loop=loop, connector=connector) + kwargs = {} + + if ws_response_class is not None: + kwargs['ws_response_class'] = ws_response_class + + session = aiohttp.ClientSession(loop=loop, connector=connector, **kwargs) try: resp = yield from session.ws_connect( url, protocols=protocols, timeout=timeout, - ws_response_class=ws_response_class, autoclose=autoclose, autoping=autoping) return resp diff --git a/docs/client_reference.rst b/docs/client_reference.rst index 0a225b79966..19bcf2cb7d6 100644 --- a/docs/client_reference.rst +++ b/docs/client_reference.rst @@ -30,8 +30,8 @@ Usage example:: .. class:: ClientSession(*, connector=None, loop=None, request_class=None,\ - response_class=None, cookies=None, headers=None,\ - auth=None) + response_class=None, ws_response_class=None,\ + cookies=None, headers=None, auth=None) The class for creating client sessions and making requests. @@ -51,6 +51,9 @@ Usage example:: :param response_class: Custom Response class implementation (optional) + :param ws_response_class: Custom WebSocketResponse class implementation + (optional) + :param dict cookies: Cookies to send with the request (optional) :param dict headers: HTTP Headers to send with @@ -225,8 +228,7 @@ Usage example:: .. coroutinemethod:: ws_connect(url, *, protocols=(), timeout=10.0\ - ws_response_class=None, autoclose=True,\ - autoping=True) + autoclose=True, autoping=True) Create a websocket connection. Returns a :class:`ClientWebSocketResponse` object. @@ -236,8 +238,6 @@ Usage example:: :param float timeout: Timeout for websocket read. 10 seconds by default - :param ws_response_class: (optional) Custom Response class implementation - :param bool autoclose: Automatically close websocket connection on close message from server. If `autoclose` is False them close procedure has to be handled manually diff --git a/docs/client_websockets.rst b/docs/client_websockets.rst index fc182c15058..4ef18e64b3b 100644 --- a/docs/client_websockets.rst +++ b/docs/client_websockets.rst @@ -62,7 +62,8 @@ do not create an instance of class :class:`ClientWebSocketResponse` manually. :param obj connector: object :class:`TCPConnector` - :param ws_response_class: (optional) Custom Response class implementation. + :param ws_response_class: (optional) Custom WebSocketResponse class + implementation. :param bool autoclose: Automatically close websocket connection on close message from server. If `autoclose` is From ae43d460bd42302ed6c8dca5dc54a94e7226bab6 Mon Sep 17 00:00:00 2001 From: davebshow Date: Thu, 21 May 2015 10:50:10 -0400 Subject: [PATCH 3/5] fixed default values for ClientSession constructor in docs, updated param order so it looks better --- aiohttp/client.py | 8 ++++---- docs/client_reference.rst | 23 ++++++++++++----------- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/aiohttp/client.py b/aiohttp/client.py index 26b1eb46b4e..b1f33196d86 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -30,10 +30,10 @@ class ClientSession: _source_traceback = None _connector = None - def __init__(self, *, connector=None, loop=None, - request_class=ClientRequest, response_class=ClientResponse, - ws_response_class=ClientWebSocketResponse, 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 diff --git a/docs/client_reference.rst b/docs/client_reference.rst index 19bcf2cb7d6..6eb79e57826 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, ws_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,14 +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 ws_response_class: Custom WebSocketResponse class implementation - (optional) - :param dict cookies: Cookies to send with the request (optional) :param dict headers: HTTP Headers to send with @@ -62,6 +54,15 @@ 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. + .. attribute:: closed ``True`` if the session has been closed, ``False`` otherwise. From 8188ebc3e797e87834f1d89038ef1077a1fe8eaf Mon Sep 17 00:00:00 2001 From: davebshow Date: Thu, 21 May 2015 21:23:50 -0400 Subject: [PATCH 4/5] added coroutine decorator, updated docs with version info --- aiohttp/websocket_client.py | 1 + docs/client_reference.rst | 17 +++++++++++++---- docs/client_websockets.rst | 2 ++ 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/aiohttp/websocket_client.py b/aiohttp/websocket_client.py index ea5841b367d..40dfb6bf3b8 100644 --- a/aiohttp/websocket_client.py +++ b/aiohttp/websocket_client.py @@ -29,6 +29,7 @@ class MsgType(IntEnum): closedMessage = Message(MsgType.closed, None, None) +@asyncio.coroutine def ws_connect(url, *, protocols=(), timeout=10.0, connector=None, ws_response_class=None, autoclose=True, autoping=True, loop=None): diff --git a/docs/client_reference.rst b/docs/client_reference.rst index 6eb79e57826..757302dc1f2 100644 --- a/docs/client_reference.rst +++ b/docs/client_reference.rst @@ -54,15 +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 request_class: Request class implementation. ``ClientRequest`` by + default. - :param response_class: Response class implementation. - ``ClientResponse`` 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. @@ -245,6 +253,7 @@ Usage example:: :param bool autoping: automatically send `pong` on `ping` message from server + .. versionadded:: 0.16 .. method:: close() diff --git a/docs/client_websockets.rst b/docs/client_websockets.rst index 4ef18e64b3b..cec20dd2153 100644 --- a/docs/client_websockets.rst +++ b/docs/client_websockets.rst @@ -65,6 +65,8 @@ do not create an instance of class :class:`ClientWebSocketResponse` manually. :param ws_response_class: (optional) Custom WebSocketResponse class implementation. + .. versionadded:: 0.16 + :param bool autoclose: Automatically close websocket connection on close message from server. If `autoclose` is False them close procedure has to be handled manually From a95360768cd3c71adedd67b453dd1a79ef72f24d Mon Sep 17 00:00:00 2001 From: davebshow Date: Thu, 21 May 2015 21:50:49 -0400 Subject: [PATCH 5/5] added ClientWebSocketResponse class as default for ws_connect --- aiohttp/websocket_client.py | 58 +++++++++++++++++-------------------- docs/client_websockets.rst | 10 +++---- 2 files changed, 32 insertions(+), 36 deletions(-) diff --git a/aiohttp/websocket_client.py b/aiohttp/websocket_client.py index 40dfb6bf3b8..718f6da4df9 100644 --- a/aiohttp/websocket_client.py +++ b/aiohttp/websocket_client.py @@ -29,37 +29,6 @@ class MsgType(IntEnum): closedMessage = Message(MsgType.closed, None, None) -@asyncio.coroutine -def ws_connect(url, *, protocols=(), timeout=10.0, connector=None, - ws_response_class=None, 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) - - kwargs = {} - - if ws_response_class is not None: - kwargs['ws_response_class'] = ws_response_class - - session = aiohttp.ClientSession(loop=loop, connector=connector, **kwargs) - - try: - resp = yield from session.ws_connect( - url, - protocols=protocols, - timeout=timeout, - autoclose=autoclose, - autoping=autoping) - return resp - - finally: - session.detach() - - class ClientWebSocketResponse: def __init__(self, reader, writer, protocol, @@ -202,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_websockets.rst b/docs/client_websockets.rst index cec20dd2153..b771c6ecc46 100644 --- a/docs/client_websockets.rst +++ b/docs/client_websockets.rst @@ -46,9 +46,9 @@ 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=(), timeout=10.0,\ - connector=None, ws_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 @@ -62,8 +62,8 @@ do not create an instance of class :class:`ClientWebSocketResponse` manually. :param obj connector: object :class:`TCPConnector` - :param ws_response_class: (optional) Custom WebSocketResponse class - implementation. + :param ws_response_class: WebSocketResponse class implementation. + ``ClientWebSocketResponse`` by default. .. versionadded:: 0.16