How to mock imwrite in a loop

Viewed 32

I am trying to use Python unittest library to test the following code.

from cv2 import imwrite
import h5py
import lumpy as np

class Myclass:
    def __init__(self):
        self.data = []

    def get_data(self):
        return self.data

    def load_data(self, path):
        count = 10
        for i in range(count):
            mat_file = path + f'{i}.mat'
            with h5py.File(mat_file, 'r') as fin:
                data = np.array(fin['data'])

            for j in range(len(data)):
                filename = path + f'{i}_{j}.jpg'
                imwrite(filename, data)

        self.data = data

I tried to do the following. But I am getting AttributeError saying Myclass has no attribute 'imwrite'. And I do not know how to mock the nested loops with image writing imwrite.

from mock import MagicMock, patch

def test():
    m = MagicMock()
    m.__enter__.return_value = data

    with patch("h5py.File", return_value=m):
        with patch(Myclass.imwrite) as mock_imwrite:
            testclass = Myclass()
            testclass.load_data(path='test')
            assert testclass.get_data() == data

I hope someone can help me out. Any help is very much appreciated

1 Answers

The issue here is that your class Myclass does not have a function imwrite. Only the module (aka the file the class is defined in) is aware of this function.

If you want to patch the imwrite function of the package cv2, you have to write:

(untested code)

with patch('cv2.imwrite') as mock_imwrite:
   testclass = Myclass()
   ...

But the as mock_imwrite part is only needed, if you run an assert on the mock. Otherwise, you may just skip it.

Related