Error in 'async with' statement when mocking coroutine that returns MagicMock

import pytest
import asyncio
from asynctest import MagicMock, CoroutineMock

class TestClass(object):
    @pytest.mark.asyncio
    async def test_get_queue_url(self):
        test_client = MagicMock()
        test_session = MagicMock()
        expected = "testvalue"

        test_client.get_test = CoroutineMock(side_effect=[expected])
        test_session.create_client = CoroutineMock(side_effect=[test_client])
        async with test_session.create_client("something") as client:
            print("hello")

For some reason, this generates the following error:

self = <test_Distributor.TestClass object at 0x00000242FDA174A8>

	@pytest.mark.asyncio
	async def test_get_queue_url(self):
		test_client = MagicMock()
		test_session = MagicMock()
		expected = "testvalue"

		test_client.get_test = CoroutineMock(side_effect=[expected])
		test_session.create_client = CoroutineMock(side_effect=[test_client])
>       async with test_session.create_client("something") as client:
E       AttributeError: __aexit__

Is this expected? I would’ve thought that the coroutine gets called and returns the MagicMock, which should be able to be used in the async with statement. Unless I’m doing something wrong, maybe someone can point me in the right direction.

Thanks!

[Martiusweb]

I think you misunderstood my comment because of the bad formatting.

When python executes

with A as B:
    pass

the interpreter calls A.aenter (a coroutine) and the result of this call is assigned to the symbol B. In you case A is test_session.create_client(), which means that the method aenter is called on the object returned by the call to test_session.create_client, not on the create_client mock itself.

this means that you need to do something like:

mock_client = MagicMock()
mock_client.get_test = CoroutineMock(return_value=expected)

context_manager = MagicMock()
context_manager.__aenter__.return_value = mock_client

mock_session = MagicMock()
mock_session.create_client.return_value = context_manager

[Martiusweb]

@Kentzo is correct but missing a small thing:

__aenter__ is called on the result of create_client(), not create_client itself, you need to do something like:

test_session.create_client().__aenter__.get_test = CoroutineMock(side_effect=[expected])

This is true as long as the create_client mock returns the same object.