Testing Async Coroutines With Context Manager

2k views Asked by At

I have been trying to test a Context managed Async coroutine through aiobotocore in Python 3.7 . I have been using the asynctest package to get the included MagicMock which has the magic methods of __aenter__ and __aexit__ and a custom mock factory that returns a MagicMock object as the result of an awaitable coroutine, but i am having trouble with the coroutine inside the context manager. The function I am trying to mock:

from aiologger import Logger
import aiobotocore


async def delete_file(bucket, key, alogger):
    await alogger.info(f'deleting file {key}')
    session = aiobotocore.get_session()
    async with session.create_client('s3', use_ssl=False) as s3:
        await s3.delete_object(
            Bucket=bucket,
            Key=key)

This is called with the input params later in the code, my testing code is:

import asyncio
from src import main
from unittest import TestCase, mock
from asynctest.mock import CoroutineMock, MagicMock as AsyncMagicMock

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

class TestMain(TestCase):

    @mock.patch('src.main.aiobotocore.get_session', new_callable=AsyncMagicMock)
    @mock.patch('src.main.Logger', new_callable=AsyncMockCall)
    def test_delete_file(self, alogger, botomock):
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main.delete_file('test_bucket',
                                                'test_key.csv',
                                                alogger))

Yet when i run it I am getting this error messasge:

____________________________________________________________________________ TestMain.test_delete_file ____________________________________________________________________________

self = <tests.test_main.TestMain testMethod=test_delete_file>, alogger = <AsyncMockCall name='Logger' id='4480486312'>, botomock = <MagicMock name='get_session' id='4480486144'>

    @mock.patch('src.main.aiobotocore.get_session', new_callable=AsyncMagicMock)
    @mock.patch('src.main.Logger', new_callable=AsyncMockCall)
    def test_delete_file(self, alogger, botomock):
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main.delete_file('test_bucket',
                                                'test_key.csv',
>                                               alogger))

tests/test_main.py:21: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/asyncio/base_events.py:584: in run_until_complete
    return future.result()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

bucket = 'test_bucket', key = 'test_key.csv', alogger = <AsyncMockCall name='Logger' id='4480486312'>

    async def delete_file(bucket, key, alogger):
        await alogger.info(f'deleting file {key}')
        session = aiobotocore.get_session()
        async with session.create_client('s3', use_ssl=False) as s3:
            await s3.delete_object(
                Bucket=bucket,
>               Key=key)
E           TypeError: object MagicMock can't be used in 'await' expression

src/main.py:20: TypeError
============================================================================= short test summary info =============================================================================
FAILED tests/test_main.py::TestMain::test_delete_file - TypeError: object MagicMock can't be used in 'await' expression

It looks to me like i need the asynctest magicmock to handle the context manager, but then i need the custom mock I made to return a coroutine. I know there is a CoroutineMock with asynctest but i cannot get it to work in this context, how would I resolve this?

1

There are 1 answers

1
JSwordy On BEST ANSWER

So from other answers it looks like i need to mock the specific delete_object method with the CoroutineMock, and my issue with getting it to work was that aiobotocore uses AioSession as the class for get session, the below code works:

import asyncio
from src import main
from unittest import TestCase, mock
from asynctest.mock import CoroutineMock, MagicMock as AsyncMagicMock

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

class TestMain(TestCase):

    @mock.patch('src.main.aiobotocore.AioSession.create_client', new_callable=AsyncMagicMock)
    @mock.patch('src.main.Logger', new_callable=AsyncMockCall)
    def test_delete_file(self, alogger, botomock):
        loop = asyncio.get_event_loop()
        botomock.return_value.__aenter__.return_value.delete_object = CoroutineMock(return_value=[])
        loop.run_until_complete(main.delete_file('test_bucket',
                                                'test_key.csv',
                                                alogger))