在Python中模拟文件对象或可迭代对象

Question

怎样用 mock 库来模拟和测试那些通过 open() 函数返回的对象的代码，才是正确的做法呢？

whitelist_data.py 文件内容如下：

WHITELIST_FILE = "testdata.txt"

format_str = lambda s: s.rstrip().lstrip('www.')
whitelist = None

with open(WHITELIST_FILE) as whitelist_data:
    whitelist = set(format_str(line) for line in whitelist_data)

if not whitelist:
    raise RuntimeError("Can't read data from %s file" % WHITELIST_FILE)

def is_whitelisted(substr):
    return 1 if format_str(substr) in whitelist else 0

这是我尝试测试的方式。

import unittest
import mock 

TEST_DATA = """
domain1.com
domain2.com
domain3.com
"""

class TestCheckerFunctions(unittest.TestCase):

    def test_is_whitelisted_method(self):
        open_mock = mock.MagicMock()
        with mock.patch('__builtin__.open',open_mock):
            manager = open_mock.return_value.__enter__.return_value
            manager.__iter__ = lambda s: iter(TEST_DATA.splitlines())
            from whitelist_data import is_whitelisted
            self.assertTrue(is_whitelisted('domain1.com'))

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

运行 python tests.py 的结果是：

$ python tests.py

E
======================================================================
ERROR: test_is_whitelisted_method (__main__.TestCheckerFunctions)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "tests.py", line 39, in test_is_whitelisted_method
    from whitelist_data import is_whitelisted
  File "/Users/supa/Devel/python/whitelist/whitelist_data.py", line 20, in <module>
    whitelist = set(format_str(line) for line in whitelist_data)
TypeError: 'Mock' object is not iterable

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

更新：感谢 Adam，我重新安装了 mock 库（pip install -e hg+https://code.google.com/p/mock#egg=mock），并更新了 tests.py。现在运行得很好。

Answer 1

你需要一个 MagicMock。这个东西可以用来进行迭代。

在 mock 0.80beta4 版本中，patch 会返回一个 MagicMock。所以这个简单的例子可以正常工作：

import mock

def foo():
    for line in open('myfile'):
        print line

@mock.patch('__builtin__.open')
def test_foo(open_mock):
    foo()
    assert open_mock.called

如果你使用的是 mock 0.7.x 版本（看起来你是这个版本），我觉得仅靠 patch 是无法实现这个功能的。你需要先单独创建这个 mock，然后再把它传入 patch：

import mock

def foo():
    for line in open('myfile'):
        print line

def test_foo():
    open_mock = mock.MagicMock()
    with mock.patch('__builtin__.open', open_mock):
        foo()
        assert open_mock.called

---

注意 - 我是用 py.test 运行这些的，不过这些方法在 unittest 中也同样适用。