Mocking zipfile in python

Question

I'm trying to use Python mock library to mock few methods for zipfile module.

Example Source which I want to test:

def zipStuff(listOfPathToFiles):
    with ZipFile(fName, 'w') as archive:
        for each in listOfPathToFiles:
             archive.write(each, strippedfName)
    return archive

The "archive" above will be ignored for normal execution, but will be list of files during tests.

Example unittest code:

emptyList=[]

def mockWrite(fName):
    emptyList.append(fName)
    return

mockZip.__enter__ = Mock(return_value=emptyList)
mockZip.__exit__ = Mock(return_value=True)

Now, I want to mock archive.write so that instead of actual write call, it is replaced by mockWrite function so I can get a list of all files that were supposed to be zipped.

I've tried:

mockZip.write = Mock(side_effect=mockWrite)

but this wasn't being called. Debugging shows the function is calling mockZip.enter().write. If I try:

mockZip.__enter__().write = Mock(side_effect=mockWrite)

Python issues error that 'list' has no attribute write (which is correct). I'm new to Mock and Python and would really appreciate any pointers. Suggestions?

Answer 1

Instead of having mockZip.__enter__ return an empty list, have it return an object like the following:

class MockZipFile:
    def __init__(self):
        self.files = []
    def __iter__(self):
        return iter(self.files)
    def write(self, fname):
        self.files.append(fname)

You can add methods and implementations as needed to suit your mocking needs.