Python: How do I mock an async method that gets called multiple times?

Viewed 5525

I have a method I use for testing async code in Python (3.6):

@asyncio.coroutine
def coroutine_creater(value):
    return value

I use it like this:

async def test_test():
    my_mock = Mock(return_value=coroutine_creater(5))

    # I call with "await"
    first_call = await my_mock()
    second_call = await my_mock()

    assert first_call == 5, 'first call failed'
    assert second_call == 5, 'second call failed'  # this call fails

This way I can create mocks for async calls. I found that this doesn't work if I call the async method twice. In my code, first_call equals 5 like I expect, but second_call equals None. What's going on here? What can I do to test code that calls a Mock async method multiple times?

1 Answers

You should set the side_effect argument for Mock.

side_effect: A function to be called whenever the Mock is called. See the side_effect attribute. Useful for raising exceptions or dynamically changing return values. The function is called with the same arguments as the mock, and unless it returns DEFAULT, the return value of this function is used as the return value.

Below example using pytest-asyncio and pytest modules:

code_53856568.py:

import asyncio


@asyncio.coroutine
def coroutine_creater(value):
    return value

test_code_53856568.py:

from unittest.mock import Mock
from code_53856568 import coroutine_creater
import pytest


@pytest.mark.asyncio
async def test_test():
    def coroutine_creater_side_effect():
        return coroutine_creater(5)

    my_mock = Mock(side_effect=coroutine_creater_side_effect)

    first_call = await my_mock()
    second_call = await my_mock()

    assert first_call == 5, 'first call failed'
    assert second_call == 5, 'second call failed'

unit test result with coverage report:

(venv) ☁  python-codelab [master] ⚡  coverage run -m pytest /Users/ldu020/workspace/github.com/mrdulin/python-codelab/src/stackoverflow/53856568/test_code_53856568.py && coverage report -m --include="src/*"
===================================================================================================================== test session starts =====================================================================================================================
platform darwin -- Python 3.7.5, pytest-5.3.1, py-1.8.0, pluggy-0.13.1
rootdir: /Users/ldu020/workspace/github.com/mrdulin/python-codelab
plugins: asyncio-0.10.0
collected 1 item                                                                                                                                                                                                                                              

src/stackoverflow/53856568/test_code_53856568.py .                                                                                                                                                                                                      [100%]

====================================================================================================================== 1 passed in 0.04s ======================================================================================================================
Name                                               Stmts   Miss  Cover   Missing
--------------------------------------------------------------------------------
src/stackoverflow/53856568/code_53856568.py            3      0   100%
src/stackoverflow/53856568/test_code_53856568.py      11      0   100%
--------------------------------------------------------------------------------
TOTAL
Related