Python unit test how to use Mox to mock the gzip with statement

Question

In Python, how do I mock an object created in a with statement using mox unit test library

Code

class MyCode:
    def generate_gzip_file(self):
        with gzip.GzipFile('file_name.txt.gz','wb') as f:
             f.write('data')

Unit Test

class MyCodeTest(unittest.TestCase):
    def test_generate_gzip_file(self):
        mox = mox.Mox()
        mock_gzip_file = self.mox.CreateMock(gzip.GzipFile)
        mox.StubOutWithMock(gzip, 'GzipFile')
        gzip.GzipFile('file_name.txt.gz','wb').AndReturn(mock_file)
        mock_gzip_file.write('data')
        mox.ReplayAll()
        MyCode().generate_gzip_file()
        mox.VerifyAll()

I get the error AttributeError: __exit__ on line

with gzip.GzipFile('file_name.txt.gz','wb') as f:

Are you using python 2.6 or lower? I don't think GzipFile grew support for context management until 2.7. You'll probably have to write your own (tiny) wrapper if so. — DSM, Jan 18 '12 at 05:16
GzipFile supports iteration and the with statement. according to documentation http://docs.python.org/library/gzip.html — venky, Jan 18 '12 at 05:45
yeah, that came in with 2.7, as I said, and since you're running 2.7 that can't be the problem. I think what's going on is that the stubbed version doesn't wind up with an `__exit__`. I think you're going to have to make sure that the with is mocked (see the examples at http://garybernhardt.github.com/python-mock-comparison/ -- search for 'context manager' and look at the Mox examples.) — DSM, Jan 18 '12 at 05:52

score 3 · Accepted Answer · answered Jan 19 '12 at 01:23

DSM is correct that the mocked instance of gzip.GzipFile isn't ending up with a __exit__ method for some reason. You'll get exactly the same error if you forget to define __exit__ on a class you use with a with statement. For example:

>>> class C(object):
...   def __enter__(self):
...     return self
... 
>>> with C() as c:
...   pass
... 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: __exit__

Fortunately, you can work around the problem by using Mox's CreateMockAnything() method to create a mock_gzip_file object that doesn't enforce a particular interface. You'll need to be careful to ensure that you set up the expectations for the mock_gzip_file object correctly (i.e. that you set up expectations for when and how the __enter__() and __exit__(...) methods will be called). Here's an example that worked for me:

import gzip
import mox
import unittest

class MyCode:
    def generate_gzip_file(self):
        with gzip.GzipFile('file_name.txt.gz', 'wb') as f:
             f.write('data')

class MyCodeTest(unittest.TestCase):
    def test_generate_gzip_file(self):
        mymox = mox.Mox()
        mock_gzip_file = mymox.CreateMockAnything()
        mymox.StubOutWithMock(gzip, 'GzipFile')
        gzip.GzipFile('file_name.txt.gz', 'wb').AndReturn(mock_gzip_file)
        mock_gzip_file.__enter__().AndReturn(mock_gzip_file)
        mock_gzip_file.write('data')
        mock_gzip_file.__exit__(None, None, None).AndReturn(None)
        mymox.ReplayAll()

        MyCode().generate_gzip_file()
        mymox.VerifyAll()

if __name__ == '__main__':
    unittest.main()

When I run this I get:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

Python unit test how to use Mox to mock the gzip with statement

1 Answers1