模拟CSV DictReader

Question

有没有一种方法可以在单元测试中模拟DictReader，而不必实际编写文件并重新打开它。

我的函数接受一个DictReader实例，所以我可以很容易地传递一个来测试它们的功能，但似乎无法在不打开文件的情况下获取一个。

目前，我正在手动编写CSV文件，然后在每个测试结束后将其删除。

class TestRowsStuff(unittest.TestCase):
    def write_csv(self, path, iterable):
        with open(path, 'wb') as f:
            writer = csv.DictWriter(f, [PP, SN, TN])
            writer.writeheader()
            writer.writerows(iterable)

    def setUp(self):
        ...
        self.test_file = os.path.join('test.csv')
        self.write_csv(self.test_file, test_values)

    def tearDown(self):
        os.remove(self.test_file)

- Jared Mackey

---

为什么不直接使用内存中的StringIO对象来存储/读取文件呢？ - Anzel

Answer 1

您可以使用内存中的StringIO对象来存储/读取Unicode/字符串：

In [10]: from StringIO import StringIO

In [11]: import csv

In [12]: csvfile = StringIO()

In [13]: csvfile.seek(0)

# sample taken from [here](https://docs.python.org/2/library/csv.html)    
In [14]: fieldnames = ['first_name', 'last_name']

In [15]: writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

In [16]: writer.writeheader()

In [17]: writer.writerow({'first_name': 'Baked', 'last_name': 'Beans'})

In [18]: writer.writerow({'first_name': 'Lovely', 'last_name': 'Spam'})

In [19]: writer.writerow({'first_name': 'Wonderful', 'last_name': 'Spam'})

In [20]: csvfile.seek(0)

要读回来：

In [21]: csvfile.readlines()
Out[21]: 
['first_name,last_name\r\n',
 'Baked,Beans\r\n',
 'Lovely,Spam\r\n',
 'Wonderful,Spam\r\n']

如果您想使用内存缓冲区，您也可以使用 io.StringIO。

- Anzel

---

这看起来很有趣。我会试一试。这不是一个经常运行的测试，只是很不方便我不得不一遍又一遍地写出文件。 - Jared Mackey

尝试时我得到了TypeError: expected str, bytes or os.PathLike object, not _io.StringIO的错误; 这个答案是针对python2编写的吗？ - coler-j

@coler-j 在 Python 3 版本中使用 from io import StringIO。 - Ali_Sh

Answer 2

如果有人正在寻找模拟writerow并想要检查数据是否已成功传递到CSV写入器的情况：

from unittest import mock

@mock.patch('csv.writer')
def test_csv_writer_writerow(self, csv_writer_mock):
    (do your stuff...)
    csv_writer_mock.writerow.assert_has_calls([mock.call(expected_data_here)])

- Degim

Answer 3

如果你正在寻找一种计算 writerow 的方法，这可能会有所帮助。

from mock import Mock, patch
from where.you.do.stuff import do
@patch('where.you.do.stuff.csv')
def test846(self, m_csv):
    m_csv.writer = Mock(writerow=Mock())
    do()
    self.assertEqual(m_csv.writer.call_count, 1)
    self.assertEqual(m_csv.writer().writerow.call_count, 999)



import csv
from io import StringIO
def do():
    w = csv.writer(StringIO)
    for i in range(1, 999):
        w.writerow(i)

- Danil