Sunday, 20 September 2015

Mocking python objects and object functions using both class-level and function-level mocks

Had some fun solving an issue with partitions larger than 2Tb, and came across a little gotcha when it comes to mocking in python when a) you want to mock both an object and a function in that object, and b) when you want to mock.patch.object at both the test class and test method level.

Say you have a function you want to test that looks like this:


def make_partitions(...):
        ...
        dp = disk_partitioner.DiskPartitioner(...)
        dp.add_partition(...)
        ...

where the DiskPartitioner class looks like this:


class DiskPartitioner(object):

    def __init__(self, ...):
        ...

    def add_partition(self, ...):
        ...


and you have existing test code like this:

@mock.patch.object(utils, 'execute')
class MakePartitionsTestCase(test_base.BaseTestCase):
    ...


and you want to add a new test function, but adding a new patch just for your new test.

You want to verify that the class is instantiated with the right options, and you need to mock the add_partition method as well. How do you use the existing test class (with the mock of the execute function), add a new mock for the DiskPartitioner.add_partition function, and the __init__ of the DiskPartitioner class?

After a little trial and error, this is how:

    @mock.patch.object(disk_partitioner, 'DiskPartitioner',
                       autospec=True)
    def test_make_partitions_with_gpt(self, mock_dp, mock_exc):

        # Need to mock the function as well
        mock_dp.add_partition = mock.Mock(return_value=None)
        ...
        disk_utils.make_partitions(...)   # Function under test
        mock_dp.assert_called_once_with(...)
        mock_dp.add_partition.assert_called_once_with(...)


Things to note:

1) The ordering of the mock parameters to test_make_partitions_with_gpt isn't immediately intuitive (at least to me).  You specify the function level mocks first, followed by the class level mocks.

2) You need to manually mock the instance method of the mocked class.  (i.e. the add_partition function)


You can see the whole enchilada over here in the review.

1 comment: