Subclassing MagicMock in python

Question

I have a class that I need to patch, which works similar to this

class Foo(object):
    def __init__(self, query):
        self._query = query

    def do_stuff(self):
        # do stuff with self._query

How would I set up a mocking class for Foo such that

foo = MockFoo(query)
foo.do_stuff()

returns a mock result, but still taking into account the data passed in for query. I thought of subclassing MagicMock like this

class MockFoo(MagicMock):
    def __init__(self, query, *args, **kwargs):
        super(MagicMock, self).__init__(*args, **kwargs)
        self._query

    def do_stuff(self):
        mock_result = self._query * 10
        return MagicMock(return_value=mock_result)

but I can't quite figure out how to apply patch to use MockFoo instead of MagicMock in the actual TestCase. I need to be able to write a test similar to this.

with patch('x.Foo') as mock_foo:
    # run test code
    self.assertTrue(mock_foo.called)
    self.assertEqual(mock_foo.wait.return_value, 20)

I know I can just use plain MagicMock with autospec=True or something and override the return values for each of the methods, but it would be nice to have a mock class which I can just use to replace the production class in one go. The key bit being, having to access the member variable self._query in the mock methods, while having it initialized in the constructor (just like the production class).

Answer 1

About you base question the answer is use new_callable parameter in patch:

with patch('x.Foo', new_callable=MockFoo) as mock_foo:
    # run test code
    self.assertTrue(mock_foo.called)
    self.assertEqual(mock_foo.wait.return_value, 20)

If you need to add some argument to MockFoo init call consider that every argument not used in patch will be passet to the mock constructor (MockFoo in your case).

If you need to a wrapper of your production class maybe you are looking in the wrong place: mock is not wrapper.

When you mock something you don't really want to know how your mock do the work but just how your code use it and how you code react to mocked objects answers.