Mocking Context Managers in Python

I've often found Python's context managers to be pretty useful. They make a nice interface that can handle starting and ending of temporary things for you, like opening and closing a file.

For example:

f = open('myfile.txt', 'w')
try:
for row in records:
f.write(row)
finally:
f.close()

can be replaced with

with open('myfile.txt', 'w') as f:
for row in records:
f.write(row)

This last week I was working with the ZipFile module and wanted to use it's context manger interface, but I ran into a little confusion when it came to unit testing. After a little better understanding of how context managers work, I figured out that the __enter__ and __exit__ methods are what really makes a context handler. As explained at PyMOTW, when you invoke with on a class, __enter__ is called and should return an object to be used in the context (f in the above example), the code within the block is executed, and __exit__ is called no matter the outcome of the block. So both of these would be roughly equivalent, assuming do_stuff doesn't raise an exception.

with Context(foo) as bar:
do_stuff(bar)
c = Context(foo)
bar = c.__enter__()
try:
do_stuff(bar)
finally:
c.__exit__(None, None, None)

With this understanding, here is the solution to my mocking problem using PyMox.

#module/tasks.py
def zip_it_up(filename):
with ZipFile(filename, 'w') as f:
for file in FILES:
f.write(file.path, file.name)
# tests.py
... # setup and stuff is up here somewhere
def test_building_zipfile(self):
self.mock.StubOutWithMock(module.tasks, 'ZipFile')
    mock_zip = self.mock.CreateMockAnything()
    module.tasks.ZipFile('/tmp/export.zip', 'w').AndReturn(mock_zip)
    mock_zip.__enter__().AndReturn(mock_zip)
    mock_zip.write('/tmp/export-1.xml', 'export-1.xml')
    mock_zip.write('/tmp/export-2.xml', 'export-2.xml')
    mock_zip.__exit__(None, None, None)
self.mock.ReplayAll()
zip_it_up()
self.mock.VerifyAll()

Mocking out ZipFile allows us to return a mock object from it's instantiation. We can then set the expectation that __enter__ will be called on the instance, returning the instance itself, expecting write to be called twice on the instance and finally __exit__ to be called. The three arguments of None here are to indicate that an exception isn't expected. If the code inside the context block were to raise an exception, these arguments would be the type, value and traceback as returned by raise. In the event you are testing for an exception, these arguments should be set accordingly when setting expectations.