diff --git a/tests/integration/aiohttp_utils.py b/tests/integration/aiohttp_utils.py index 2f135c8..e9f9364 100644 --- a/tests/integration/aiohttp_utils.py +++ b/tests/integration/aiohttp_utils.py @@ -1,13 +1,22 @@ +import asyncio + import aiohttp -async def aiohttp_request(loop, method, url, output='text', **kwargs): # NOQA: E999 - async with aiohttp.ClientSession(loop=loop) as session: # NOQA: E999 - async with session.request(method, url, **kwargs) as response: # NOQA: E999 - if output == 'text': - content = await response.text() # NOQA: E999 - elif output == 'json': - content = await response.json() # NOQA: E999 - elif output == 'raw': - content = await response.read() # NOQA: E999 - return response, content +@asyncio.coroutine +def aiohttp_request(loop, method, url, output='text', **kwargs): + session = aiohttp.ClientSession(loop=loop) + response_ctx = session.request(method, url, **kwargs) # NOQA: E999 + + response = yield from response_ctx.__aenter__() # NOQA: E999 + if output == 'text': + content = yield from response.text() # NOQA: E999 + elif output == 'json': + content = yield from response.json() # NOQA: E999 + elif output == 'raw': + content = yield from response.read() # NOQA: E999 + + response_ctx._resp.close() + yield from session.close() + + return response, content