QueryMock

querymock.py

class QueryMock(object):
    """ Beats having to mock query.return_all.filter...

    Usage:
    >>> qm = QueryMock([("all", "ERROR")])
    >>> getDatabaseSession = MagicMock(return_value=qm)
    >>> errors = (
    ...     getDatabaseSession()
    ...     .query(ShowTime)
    ...     .filter(ShowTime.start >= start)
    ...     .filter(ShowTime.end < end)
    ...     .all()
    ... )
    >>> print errors
    ERROR
    """
    def __init__(self, func_mock):
        """
        :param func_mock: Iterable of tuples, where the first element
            is the name of the function name that should return the value
            of the second element. See the example in the class doc for more
            information.
        """
        self.func_mock = func_mock

    def __getattribute__(self, item):
        if "func_mock" == item:
            return object.__getattribute__(self, item)
        for (func, func_return) in self.func_mock:
            if item == func:
                return lambda: func_return
        else:
            return self

    def __call__(self, *args, **kwargs):
        return self