gistfile1.py

from mock import patch
from contextlib import contextmanager


def _patch_method(obj, method_name, return_value=DEFAULT, **kwargs):
    # Allow specifying a dotted string instead of an object
    if isinstance(obj, basestring):
        return patch('%s.%s' % (obj, method_name), autospec=True, return_value=return_value, **kwargs)

    kwargs['autospec'] = True

    assert callable(getattr(obj, method_name))

    if is_exception(return_value):
        return patch.object(obj, method_name,
                            side_effect=return_value,
                            **kwargs)
    else:
        return patch.object(obj, method_name,
                            return_value=return_value,
                            **kwargs)


@contextmanager
def patch_methods(obj, *patch_method_args, **kwargs):

    require_calls = kwargs.pop('_require_calls_', True)

    patchers = [_patch_method(obj, k, v) for k, v in kwargs.items()]

    try:
        # Keep a list of the patchers that has had .start() *successfully* called
        started_patchers = []
        mocks = []

        for patcher in patchers:
            mocks.append(patcher.start())  # Might raise exception on non-existent method
            started_patchers.append(patcher)

        yield

    finally:
        # Clean up patchers that was actually applied
        for patcher in started_patchers:
            patcher.stop()

    if require_calls:
        # This code will only be executed if all mocks can be set up and
        # the block does not raise any exception
        for mock in mocks:
            assert mock.mock_calls, '%s has no calls!' % repr(mock)