WIP fix#2525

aio-libs · Nov 17, 2017 · 0969027 · 0969027
1 parent d6426ad
commit 0969027
Showing 1 changed file with 6 additions and 2 deletions.
diff --git 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,12 +278,16 @@ 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.
 
         """
+        return _WSRequestContextManager(
+            self._ws_connect(path), *args, **kwargs)
+
+    async def _ws_connect(self, path, *args, **kwargs):
         ws = await self._session.ws_connect(
             self.make_url(path), *args, **kwargs)
         self._websockets.append(ws)