为迭代自定义unittest.mock.mock_open

Question

我应该怎么调整unittest.mock.mock_open来处理这段代码呢？

file: impexpdemo.py
def import_register(register_fn):
    with open(register_fn) as f:
        return [line for line in f]

我第一次尝试使用了 read_data。

class TestByteOrderMark1(unittest.TestCase):
    REGISTER_FN = 'test_dummy_path'
    TEST_TEXT = ['test text 1\n', 'test text 2\n']

    def test_byte_order_mark_absent(self):
        m = unittest.mock.mock_open(read_data=self.TEST_TEXT)
        with unittest.mock.patch('builtins.open', m):
            result = impexpdemo.import_register(self.REGISTER_FN)
            self.assertEqual(result, self.TEST_TEXT)

但是这次尝试失败了，可能是因为代码没有使用read、readline或readlines这些方法。unittest.mock.mock_open的文档中提到，“read_data是一个字符串，用于文件句柄的read()、readline()和readlines()方法返回。当调用这些方法时，会从read_data中读取数据，直到数据用完。对这些方法的模拟相对简单。如果你需要更好地控制传给被测试代码的数据，你就需要自己定制这个模拟。默认情况下，read_data是一个空字符串。”

由于文档没有给出需要什么样的定制提示，我又尝试了 return_value 和 side_effect，但都没有成功。

class TestByteOrderMark2(unittest.TestCase):
    REGISTER_FN = 'test_dummy_path'
    TEST_TEXT = ['test text 1\n', 'test text 2\n']

    def test_byte_order_mark_absent(self):
        m = unittest.mock.mock_open()
        m().side_effect = self.TEST_TEXT
        with unittest.mock.patch('builtins.open', m):
            result = impexpdemo.import_register(self.REGISTER_FN)
            self.assertEqual(result, self.TEST_TEXT)

Answer 1

我找到了解决方案：

text_file_data = '\n'.join(["a line here", "the second line", "another 
line in the file"])
with patch('__builtin__.open', mock_open(read_data=text_file_data), 
create=True) as m:
    # mock_open doesn't properly handle iterating over the open file with for line in file:
    # but if we set the return value like this, it works.
    m.return_value.__iter__.return_value = text_file_data.splitlines()
    with open('filename', 'rU') as f:
        for line in f:
            print line

Answer 2

从Python 3.6开始，使用unittest.mock_open方法返回的模拟文件对象不支持迭代。这个问题在2014年就被报告了，到2017年时仍然没有解决。

因此，像下面这样的代码在运行时不会产生任何迭代结果：

f_open = unittest.mock.mock_open(read_data='foo\nbar\n')
f = f_open('blah')
for line in f:
  print(line)

你可以通过给模拟对象添加一个方法，来解决这个限制，使其能够正确返回行迭代器：

def mock_open(*args, **kargs):
  f_open = unittest.mock.mock_open(*args, **kargs)
  f_open.return_value.__iter__ = lambda self : iter(self.readline, '')
  return f_open

Answer 3

mock_open()这个对象确实不支持迭代。

如果你不把文件对象当作上下文管理器使用，可以试试：

m = unittest.mock.MagicMock(name='open', spec=open)
m.return_value = iter(self.TEST_TEXT)

with unittest.mock.patch('builtins.open', m):

这样一来，open()就会返回一个迭代器，这种东西可以像文件对象一样直接进行迭代，也可以和next()一起使用。不过，它不能作为上下文管理器使用。

你可以把这个和mock_open()结合起来，然后在返回值上提供__iter__和__next__方法，这样就能享受到mock_open()作为上下文管理器使用的好处：

# Note: read_data must be a string!
m = unittest.mock.mock_open(read_data=''.join(self.TEST_TEXT))
m.return_value.__iter__ = lambda self: self
m.return_value.__next__ = lambda self: next(iter(self.readline, ''))

这里的返回值是一个MagicMock对象，它是根据file对象（Python 2）或者内存文件对象（Python 3）来创建的，但只有read、write和__enter__方法被简单实现了。

上述方法在Python 2中不适用，因为a) Python 2期望存在next而不是__next__，b) 在Mock中next不被视为特殊方法（这是对的），所以即使你把上面例子中的__next__改名为next，返回值的类型也不会有next方法。对于大多数情况，只需要让生成的文件对象是一个可迭代对象而不是迭代器，可以用：

# Python 2!
m = mock.mock_open(read_data=''.join(self.TEST_TEXT))
m.return_value.__iter__ = lambda self: iter(self.readline, '')

任何使用iter(fileobj)的代码都会正常工作（包括for循环）。

在Python追踪器中有一个未解决的问题，旨在解决这个缺陷。