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.
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
__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
traceback as returned by
raise. In the event you are testing for an exception, these arguments should be set accordingly when setting expectations.