如何使用Python的'unittest'进行文件写入函数的单元测试

Question

我有一个Python函数，它会把输出文件写到磁盘上。

我想用Python的unittest模块为这个函数写一个单元测试。

我应该怎么检查文件是否相等呢？如果文件内容和预期的不一样，我希望能得到一个错误提示，并且列出不同之处。就像Unix的diff命令的输出那样。

有没有官方或者推荐的做法呢？

Answer 1

我总是尽量避免把文件写入磁盘，即使是写到一个专门用于测试的临时文件夹里。因为不实际去操作磁盘会让你的测试速度快很多，特别是当你的代码中频繁与文件交互时。

假设你有一段“很棒”的软件，保存在一个叫做 main.py 的文件里：

"""
main.py
"""

def write_to_file(text):
    with open("output.txt", "w") as h:
        h.write(text)

if __name__ == "__main__":
    write_to_file("Every great dream begins with a dreamer.")

要测试 write_to_file 这个方法，你可以在同一个文件夹里写一个叫 test_main.py 的文件，内容可以是这样的：

"""
test_main.py
"""
from unittest.mock import patch, mock_open

import main


def test_do_stuff_with_file():
    open_mock = mock_open()
    with patch("main.open", open_mock, create=True):
        main.write_to_file("test-data")

    open_mock.assert_called_with("output.txt", "w")
    open_mock.return_value.write.assert_called_once_with("test-data")

Answer 2

我更喜欢让输出函数直接接受一个文件的句柄（或者类似文件的对象），而不是接受一个文件名称并自己去打开文件。这样，我就可以在单元测试中把一个StringIO对象传给输出函数，然后通过.read()从这个StringIO对象中读取内容（在调用.seek(0)之后），并与我预期的输出进行比较。

举个例子，我们可以把这样的代码

##File:lamb.py
import sys


def write_lamb(outfile_path):
    with open(outfile_path, 'w') as outfile:
        outfile.write("Mary had a little lamb.\n")


if __name__ == '__main__':
    write_lamb(sys.argv[1])



##File test_lamb.py
import unittest
import tempfile

import lamb


class LambTests(unittest.TestCase):
    def test_lamb_output(self):
        outfile_path = tempfile.mkstemp()[1]
        try:
            lamb.write_lamb(outfile_path)
            contents = open(tempfile_path).read()
        finally:
            # NOTE: To retain the tempfile if the test fails, remove
            # the try-finally clauses
            os.remove(outfile_path)
        self.assertEqual(contents, "Mary had a little lamb.\n")

改成这样的代码

##File:lamb.py
import sys


def write_lamb(outfile):
    outfile.write("Mary had a little lamb.\n")


if __name__ == '__main__':
    with open(sys.argv[1], 'w') as outfile:
        write_lamb(outfile)



##File test_lamb.py
import unittest
from io import StringIO

import lamb


class LambTests(unittest.TestCase):
    def test_lamb_output(self):
        outfile = StringIO()
        # NOTE: Alternatively, for Python 2.6+, you can use
        # tempfile.SpooledTemporaryFile, e.g.,
        #outfile = tempfile.SpooledTemporaryFile(10 ** 9)
        lamb.write_lamb(outfile)
        outfile.seek(0)
        content = outfile.read()
        self.assertEqual(content, "Mary had a little lamb.\n")

这种方法还有一个好处，就是让你的输出函数更灵活，比如说，如果你决定不想写入文件，而是写入其他的缓冲区，因为它可以接受所有类似文件的对象。

需要注意的是，使用StringIO是基于测试输出的内容可以放进主内存。如果输出非常大，你可以使用临时文件的方法（例如，tempfile.SpooledTemporaryFile）。

Answer 3

最简单的方法是先把输出文件写出来，然后读取它的内容，再读取预期的文件内容，最后用简单的字符串比较来看看它们是否相同。如果相同，就删除输出文件；如果不同，就抛出一个错误。

这样，当测试完成后，每个失败的测试都会有一个输出文件，你可以用一些第三方工具来对比这些文件和预期文件（比如Beyond Compare就非常好用）。

如果你真的想自己提供对比的输出，记得Python的标准库里有一个叫做difflib的模块。Python 3.1的新单元测试支持中包含了一个assertMultiLineEqual方法，它使用这个模块来显示差异，效果类似于这样：

    def assertMultiLineEqual(self, first, second, msg=None):
        """Assert that two multi-line strings are equal.

        If they aren't, show a nice diff.

        """
        self.assertTrue(isinstance(first, str),
                'First argument is not a string')
        self.assertTrue(isinstance(second, str),
                'Second argument is not a string')

        if first != second:
            message = ''.join(difflib.ndiff(first.splitlines(True),
                                                second.splitlines(True)))
            if msg:
                message += " : " + msg
            self.fail("Multi-line strings are unequal:\n" + message)