gistfile1.py

import mock

class GeneratorMock(mock.Mock):
    def __init__(self, *arg, **kw):
        kw['side_effect'] = self._side_effect
        super(GeneratorMock, self).__init__(*arg, **kw)
        self._lookup = {}

    def _side_effect(self, *arg, **kw):
        if self._mock_return_value is not mock.sentinel.DEFAULT:
            return self._mock_return_value
        else:
            return self._generate(*arg, **kw)

    def _generate(self, *arg, **kw):
        key = arg + tuple((k, kw[k]) for k in sorted(kw))
        if key in self._lookup:
            return self._lookup[key]
        else:
            self._lookup[key] = ret = GeneratorMock()
            return ret

sess = GeneratorMock()

sess.query().filter(x=5).first.return_value = 5
sess.query().filter(x=2).first.return_value = 2

print(sess.query().filter(x=5).first())  # returns 5
print(sess.query().filter(x=2).first())  # returns 2
print(sess.query().filter(x=9).first())  # returns a GeneratorMock