74

How do I mock async call from one native coroutine to other one using unittest.mock.patch?

I currently have quite an awkward solution:

class CoroutineMock(MagicMock):
    def __await__(self, *args, **kwargs):
        future = Future()
        future.set_result(self)
        result = yield from future
        return result

Then

class TestCoroutines(TestCase):
    @patch('some.path', new_callable=CoroutineMock)
    def test(self, mock):
        some_action()
        mock.assert_called_with(1,2,3)

This works but looks ugly. Is there more pythonic way to do this?

Zozz
  • 1,875
  • 1
  • 14
  • 14
  • 1
    Also this mock doesn't work with asyncio.await because of asyncio.tasks.ensure_future – Zozz Sep 09 '15 at 14:33

9 Answers9

85

The solution was actually quite simple: I just needed to convert __call__ method of mock into coroutine:

class AsyncMock(MagicMock):
    async def __call__(self, *args, **kwargs):
        return super(AsyncMock, self).__call__(*args, **kwargs)

This works perfectly, when mock is called, code receives native coroutine

Example usage:

@mock.patch('my.path.asyncio.sleep', new_callable=AsyncMock)
def test_stuff(sleep):
    # code
KevinG
  • 882
  • 3
  • 16
  • 33
Zozz
  • 1,875
  • 1
  • 14
  • 14
  • 6
    This is nice, but it doesn't play well with autospec, which is basically mandatory when using MagicMock. Any thoughts on how to get that working? I'm not familiar enough with the internals... – Symmetric Oct 28 '15 at 22:58
  • 4
    It works perfectly for me. I used it like this: ``` @mock.patch( 'my.path.asyncio.sleep', new_callable=AsyncMock, ) def test_stuff(sleep): # code ``` – karantan Feb 23 '17 at 08:58
  • 1
    This works. I initially like the other solution below by Ivan Castellanos. But the future is never executed and I tried like hell, but didn't get it to work. – Ytsen de Boer Oct 19 '18 at 12:48
  • 2
    This is as handy and elegant as it gets. At least for basic usage. – Felix Feb 01 '20 at 12:53
  • 11
    Note: `AsyncMock` is available in `unittest.mock` as of python 3.8, mock also automatically detects where it should be used (see [docs for examples](https://docs.python.org/3/library/unittest.mock.html#unittest.mock.AsyncMock)). – tutuDajuju Mar 18 '21 at 11:29
  • 1
    if you use python version less than 3.8, you can use `pip install mock` to get AsyncMock, https://github.com/testing-cabal/mock – Hohenheim Jun 15 '21 at 10:43
  • @tutuDajuju you should post that as an answer, along with an example. I would +1 – Jerther May 18 '22 at 13:38
  • how you write a side_effect with asyncmock ? – Valentin Garreau Sep 20 '22 at 16:04
83

Everyone's missing what's probably the simplest and clearest solution:

@patch('some.path')
def test(self, mock):
    f = asyncio.Future()
    f.set_result('whatever result you want')
    process_smtp_message.return_value = f
    mock.assert_called_with(1, 2, 3)

remember a coroutine can be thought of as just a function which is guaranteed to return a future which can, in turn be awaited.

SColvin
  • 11,584
  • 6
  • 57
  • 71
  • 8
    what is process_smtp_message.return_value = f? Also, where is the call to the function being tested? – Skorpeo Feb 23 '18 at 14:12
  • 2
    @Skorpeo - i think he means mock.return_value = f – Richard Jun 28 '18 at 08:47
  • I definitely don't. Mock is the test fixture. process_smtp_message is whatever you're trying to mock obviously. – SColvin Jun 28 '18 at 09:27
  • 15
    The solution is not compatible with Python 3.8 where is a native `AsyncMock` was introduced. So a code with the solution will fail with errors because of Future class issues. But a solution of Zozz with a simple `AsyncMock` implementation can work in Python 3.7 (for example) and Python 3.8 (if you will do conditional importing of native `AsyncMock`) – Sergey Nevmerzhitsky Jan 06 '20 at 10:08
  • For ```python3.8``` and above I ended up using: ```patch.object(your_object, 'async_method_name', return_value=your_return_value)``` – Kishan Kishore Jun 22 '21 at 12:15
  • re: "... obviously." If it was obvious, multiple people wouldn't have been confused by it. – pdoherty926 Jun 01 '22 at 14:05
28

Based on @scolvin answer I created this (imo) cleaner way:

def async_return(result):
    f = asyncio.Future()
    f.set_result(result)
    return f

That's it, just use it around whatever return you want to be async, as in

mock = MagicMock(return_value=async_return("Example return"))
await mock()
Ivan Castellanos
  • 8,041
  • 1
  • 47
  • 42
16

Subclassing MagicMock will propagate your custom class for all the mocks generated from your coroutine mock. For instance, AsyncMock().__str__ will also become an AsyncMock which is probably not what you're looking for.

Instead, you might want to define a factory that creates a Mock (or a MagicMock) with custom arguments, for instance side_effect=coroutine(coro). Also, it might be a good idea to separate the coroutine function from the coroutine (as explained in the documentation).

Here is what I came up with:

from asyncio import coroutine

def CoroMock():
    coro = Mock(name="CoroutineResult")
    corofunc = Mock(name="CoroutineFunction", side_effect=coroutine(coro))
    corofunc.coro = coro
    return corofunc

An explanation of the different objects:

  • corofunc: the coroutine function mock
  • corofunc.side_effect(): the coroutine, generated for each call
  • corofunc.coro: the mock used by the coroutine to get the result
  • corofunc.coro.return_value: the value returned by the coroutine
  • corofunc.coro.side_effect: might be used to raise an exception

Example:

async def coro(a, b):
    return await sleep(1, result=a+b)

def some_action(a, b):
    return get_event_loop().run_until_complete(coro(a, b))

@patch('__main__.coro', new_callable=CoroMock)
def test(corofunc):
    a, b, c = 1, 2, 3
    corofunc.coro.return_value = c
    result = some_action(a, b)
    corofunc.assert_called_with(a, b)
    assert result == c
Vincent
  • 12,919
  • 1
  • 42
  • 64
  • 1
    this doesn't work, side_effect=coroutine(coro), coroutine is not defined – Skorpeo Feb 24 '18 at 20:20
  • I actually like the original solution quite a bit better, as it doesn't require special rewriting of the test function. Any advantages to this approach over the one shown in the question? – Mack Sep 04 '20 at 19:20
7

Another way of mocking coroutine is to make coroutine, that returns mock. This way you can mock coroutines that will be passed into asyncio.wait or asyncio.wait_for.

This makes more universal coroutines though makes setup of tests more cumbersome:

def make_coroutine(mock)
    async def coroutine(*args, **kwargs):
        return mock(*args, **kwargs)
    return coroutine


class Test(TestCase):
    def setUp(self):
        self.coroutine_mock = Mock()
        self.patcher = patch('some.coroutine',
                             new=make_coroutine(self.coroutine_mock))
        self.patcher.start()

    def tearDown(self):
        self.patcher.stop()
Zozz
  • 1,875
  • 1
  • 14
  • 14
7

I don't know why nobody mentioned the default option available. python provides a Async version of MagicMock.

You can read more about this here. https://docs.python.org/3/library/unittest.mock.html#unittest.mock.AsyncMock

In case you are using patch then also you don't need to make any other change. It will automatically replace it with a async mock function if required. Read more here https://docs.python.org/3/library/unittest.mock.html#patch

targhs
  • 1,477
  • 2
  • 16
  • 29
  • 8
    Perhaps because the thread refers to Python 3.5 (and, by extension, all subsequent versions prior to 3.8, when `AsyncMock` was introduced) – Peter Halverson Mar 15 '22 at 16:01
  • 1
    Question was asked in 2015 when Python 3.5 was probably the "newest Python". Now in 2023 this isn't true and most people who get here from Google look for solution for much more fresh Python. So this answer is really good for today. – Konstantin Smolyanin Feb 04 '23 at 17:26
6

One more variant of "simplest" solution to mock a async object, which is just a one liner.

In source:

class Yo:
    async def foo(self):
        await self.bar()
    async def bar(self):
        # Some code

In test:

from asyncio import coroutine

yo = Yo()
# Here bounded method bar is mocked and will return a customised result.
yo.bar = Mock(side_effect=coroutine(lambda:'the awaitable should return this'))
event_loop.run_until_complete(yo.foo())
Tim
  • 2,510
  • 1
  • 22
  • 26
Murphy Meng
  • 237
  • 3
  • 8
1

You can set the return_value of an async method like so:

mock = unittest.mock.MagicMock()
mock.your_async_method.return_value = task_from_result(your_return_value)

async def task_from_result(result):
    return result

The caller will have to do await your_async_method(..) just like as if the method wasn't mocked.

JBSnorro
  • 6,048
  • 3
  • 41
  • 62
1

I like this approach, which also makes AsyncMock to behave exactly like Mock:

class AsyncMock:
    def __init__(self, *args, **kwargs):
        self.mock = Mock(*args, **kwargs)

    async def __call__(self, *args, **kwargs):
        return self.mock(*args, **kwargs)

    def __getattr__(self, item):
        return getattr(self.mock, item)

Then you can work with it in the same way like with Mock, i.e:

@pytest.mark.asyncio
async def test_async_mock_example(monkeypatch):
    fn = AsyncMock(side_effect=ValueError)
    with pytest.raises(ValueError):
        await fn()
    assert fn.call_count == 1
zhukovgreen
  • 1,551
  • 16
  • 26