Python 单元测试如何使用 Mox 模拟 gzip with 语句

Question

在Python中，我该如何使用mox单元测试库来模拟一个在with语句中创建的对象

代码

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):
        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()

我在这一行遇到了错误 AttributeError: __exit__

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

Answer 1

DSM说得对，模拟的 gzip.GzipFile 实例没有正确地包含 __exit__ 方法，这可能是出了一些问题。如果你在使用 with 语句的类中忘记定义 __exit__ 方法，你也会遇到同样的错误。举个例子：

>>> 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__

幸运的是，你可以通过使用 Mox 的 CreateMockAnything() 方法来解决这个问题，这样可以创建一个不强制要求特定接口的 mock_gzip_file 对象。不过，你需要小心确保正确设置 mock_gzip_file 对象的期望值（也就是说，你需要设置好 __enter__() 和 __exit__(...) 方法何时以及如何被调用的期望）。下面是一个对我有效的例子：

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()

当我运行这个时，我得到了：

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

OK