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

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:

Solution

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