How to mock "async with" statements?

Question

I'm trying to write tests for a method that uses "async with" statements (in this case, aioredis's connection pool), i want to mock the connection to redis, but i'm having trouble figuring out how.

Here's what i have so far:

from asyncio import Future
from unittest.mock import MagicMock

import pytest

# The thing i'm trying to test
async def set_value(redis, value):
    # Do things
    async with redis.get() as conn:
        await conn.set("key", value)

#My Mock classes
class MockRedis():
    def get(self):
        return MockAsyncPool()


class MockAsyncPool(MagicMock):
    async def __aenter__(self):
        conn = MagicMock()
        f = Future()
        f.set_result(True)
        conn.set = MagicMock(return_value=f)
        return conn

    def __aexit__(self, exc_type, exc_val, exc_tb):
        pass


# The actual test
@pytest.mark.asyncio
async def test_get_token():
    redis = MockRedis()

    token = await set_value(redis, 'something')
    assert token is not None

I run it with:

py.test path/to/file.py

And i'm getting this error:

> await conn.set("key", value)

E TypeError: object NoneType can't be used in 'await' expression

Answer 1

__aexit__ needs to also be asyncronous (needs to return an awaitable):

async def __aexit__(self, exc_type, exc_val, exc_tb):
    pass

Without being async it is returning None instead of a coroutine so it raises an error, as for the very misleading error message I have created this issue to point out that the error message needs to be fixed.