def inner(return_value=None, side_effect=None): stub = stub_connection(return_value=return_value, side_effect=side_effect) context_manager = MagicMock() context_manager.__aenter__ = CoroutineMock(return_value=stub) context_manager.__aexit__ = CoroutineMock() manager = MagicMock() manager.new_connection = Mock(return_value=context_manager) manager.connection_stub = stub return manager