From 278d5d57361b17edc5069bfc87cc87eca7680eba Mon Sep 17 00:00:00 2001 From: Christian Barra Date: Sat, 18 Nov 2017 12:16:47 +0100 Subject: [PATCH] fix async with in TestClient.ws_connect (#2528) * WIP fix#2525 * Update test_utils.py --- aiohttp/test_utils.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/aiohttp/test_utils.py b/aiohttp/test_utils.py index cd8cf7aa0f1..04ee99eb891 100644 --- a/aiohttp/test_utils.py +++ b/aiohttp/test_utils.py @@ -15,7 +15,7 @@ from yarl import URL import aiohttp -from aiohttp.client import _RequestContextManager +from aiohttp.client import _RequestContextManager, _WSRequestContextManager from . import ClientSession, hdrs from .helpers import sentinel @@ -278,13 +278,19 @@ def delete(self, path, *args, **kwargs): self.request(hdrs.METH_DELETE, path, *args, **kwargs) ) - async def ws_connect(self, path, *args, **kwargs): + def ws_connect(self, path, *args, **kwargs): """Initiate websocket connection. The api corresponds to aiohttp.ClientSession.ws_connect. """ - ws = await self._session.ws_connect( + return _WSRequestContextManager( + self._ws_connect(path, *args, **kwargs) + ) + + @asyncio.coroutine + def _ws_connect(self, path, *args, **kwargs): + ws = yield from self._session.ws_connect( self.make_url(path), *args, **kwargs) self._websockets.append(ws) return ws