From d3902a46d1691aaa074c22afadaa2da39aab4ab9 Mon Sep 17 00:00:00 2001 From: Aurelien Busi Date: Thu, 24 Sep 2020 09:40:37 +0200 Subject: [PATCH] Use context_factory as an async ctx manager like for a normal query --- tartiflette_aiohttp/_subscription_ws_handler.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tartiflette_aiohttp/_subscription_ws_handler.py b/tartiflette_aiohttp/_subscription_ws_handler.py index 12bd66c..4dfac5d 100644 --- a/tartiflette_aiohttp/_subscription_ws_handler.py +++ b/tartiflette_aiohttp/_subscription_ws_handler.py @@ -262,7 +262,8 @@ async def _handle_request(self) -> None: async def __call__(self, request: "Request") -> "WebSocketResponse": self._socket = web.WebSocketResponse(protocols=(WS_PROTOCOL,)) - self._context = await self._context_factory(request) - await self._socket.prepare(request) - await shield(self._handle_request()) - return self._socket + async with self._context_factory(request) as ctx: + self._context = ctx + await self._socket.prepare(request) + await shield(self._handle_request()) + return self._socket