Skip to content

support new param 'response_class' in websocket_client.ws_connect. #367

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
asvetlov merged 5 commits into from
May 17, 2015
Merged

support new param 'response_class' in websocket_client.ws_connect. #367

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d284ae3
support new param for user implemented response with websocket_clien…
davebshow May 17, 2015
840fbb4
added param in docs
davebshow May 17, 2015
e507683
using 'is None' idiom to set ClientWebSocketResponse class
davebshow May 17, 2015
f2180a3
updated docs to reflect order of kwargs in ws_connect function
davebshow May 17, 2015
d742d11
updated docs to reflect order of kwargs in ws_connect function
davebshow May 17, 2015
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions aiohttp/websocket_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ class MsgType(IntEnum):

@asyncio.coroutine
def ws_connect(url, protocols=(), timeout=10.0, connector=None,
autoclose=True, autoping=True, loop=None):
response_class=None, autoclose=True, autoping=True, loop=None):
"""Initiate websocket connection."""
if loop is None:
loop = asyncio.get_event_loop()
Expand Down Expand Up @@ -87,7 +87,9 @@ def ws_connect(url, protocols=(), timeout=10.0, connector=None,
reader = resp.connection.reader.set_parser(WebSocketParser)
writer = WebSocketWriter(resp.connection.writer, use_mask=True)

return ClientWebSocketResponse(
response_class = response_class or ClientWebSocketResponse
Copy link
Member

@asvetlov asvetlov May 17, 2015

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please use

if response_class is None:
    response_class = ClientWebSocketResponse idiom


return response_class(
reader, writer, protocol, resp, timeout, autoclose, autoping, loop)


Expand Down
4 changes: 3 additions & 1 deletion docs/client_websockets.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ 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, autoclose=True, autoping=True, loop=None)
.. py:function:: ws_connect(url, protocols=(), connector=None, autoclose=True, autoping=True, response_class=None, loop=None)

This function creates a websocket connection, checks the response and
returns a :class:`ClientWebSocketResponse` object. In case of failure
Expand All @@ -64,6 +64,8 @@ do not create an instance of class :class:`ClientWebSocketResponse` manually.

:param bool autoping: automatically send `pong` on `ping` message from server

:param response_class: (optional) Custom Response class implementation.

:param loop: :ref:`event loop<asyncio-event-loop>` used
for processing HTTP requests.

Expand Down
27 changes: 27 additions & 0 deletions tests/test_websocket_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,33 @@ 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):

class CustomResponse(websocket_client.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: 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)

res = self.loop.run_until_complete(
websocket_client.ws_connect(
'http://test.org',
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):
Expand Down