Mocking async call in python 3.5

Question

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?

Answer 1

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