如何最佳地对处理 CSV 文件的函数进行单元测试？

Question

我正在尝试使用Pytest对处理csv文件的函数进行单元测试。虽然我的函数可以正常工作，但在创建“样本”csv文件以测试函数时，感觉代码重复很多，实际包含真实数据的csv文件有数百万条记录。

这些不是我要在模块中测试的唯一csv文件，因此，了解处理不同文件结构的函数的最佳测试方法将非常有帮助。

目前，我正在创建一个非常简短的csv文件，以模仿实际文件模式及其单行数据，然后通过函数处理后预期生成数据框输出。

也许利用 mocking 是可行的？但是我认为你不应该需要为这种类型的测试进行 mocking。

测试函数

@pytest.mark.parametrize('test_file, expected', [
    (r'Path\To\Project\Output\Folder\mock_sales1.csv',
     pd.DataFrame([['A0A0A0', 1, 4000]], columns=['Postal_Code', 'Store_Num', 'Sales'])),
    (r'Path\To\Project\Output\Folder\mock_sales2.csv',
     pd.DataFrame([['A0A0A0', 1, 4000]], columns=['Postal_Code', 'Store_Num', 'Sales']))
])
def test_sales_dataframe(test_file, expected):
    # This part is repetitive, different tests each need a seperate file written within the test function.
    # Writing sample file to test that files with 7 columns are read correctly.
    mock_mks_sales1 = [['Data0', 'A0A0A0', 1, 'Data3', 'Data4', 'Data5', 4000]]
    with open(r'Path\To\Project\Output\Folder\mock_sales1.csv', 'w') as file:
        writer = csv.writer(file)
        writer.writerows(mock_sales1)
    # Writing sample file to test that files with 8 columns are read correctly.
    mock_mks_sales2 = [['Data0', 'A0A0A0', 1, 'Data3', 'Data4', 'Data5', 'Data6', 4000]]
    with open(r'Path\To\Project\Output\Folder\mock_sales2.csv', 'w') as file:
        writer = csv.writer(file)
        writer.writerows(mock_sales2)

    sales_df = mks_sales_dataframe(test_file)
    testing.assert_frame_equal(expected, sales_df)

    os.remove(r'Path\To\Project\Output\Folder\mock_sales1.csv')
    os.remove(r'Path\To\Project\Output\Folder\mock_sales2.csv')

主函数

def sales_dataframe(file):
    try:
        with open(file, 'r') as f:
            reader = csv.reader(f)
            num_cols = len(next(reader))
            columns = [1, 2, (num_cols - 1)]  # Number of columns is variable, this is used later to accurately specify which columns should be read. This is part I'm testing!

        sales_df = pd.read_csv(file, usecols=columns, names=['Postal_Code', 'Store_Num', 'Sales'])
        return sales_df
    except FileNotFoundError:
        raise FileNotFoundError(file)

测试按预期通过。但是，每当我要进行不同的测试时，都必须在测试函数中创建一个样本csv文件，并在测试完成后删除每个文件。可以想象，这在单个测试函数中会产生大量重复的代码，感觉相当笨重和啰嗦，特别是当测试被参数化时。

- ShockDoctor

---

我理解的对吗？在 test_sales_dataframe 中，您首先创建了 mock_sales1.csv 和 mock_sales2.csv 文件，并填充了固定内容，然后调用 mks_sales_dataframe 读取其中一个 .csv 文件，最后检查结果是否等于 expected？ - pschill

你可以将 mock_sales1.csv 和 mock_sales2.csv（以及所有其他测试数据文件）存储在与测试代码相邻的 testdata 文件夹中。然后，你的测试就变成了两行代码：sales_df = mks_sales_dataframe(test_file) 和 testing.assert_frame_equal(expected, sales_df)。 - pschill

@pschill 您的第一条评论是正确的。 - ShockDoctor

@pschill那是我最初的做法，但我认为在每次使用后删除文件可以减少混乱。最终，会有越来越多的文件（我必须读取和测试许多不同来源的数据），因此我觉得创建一个永久存储测试文件的目录是不必要的，直到实际运行测试为止。基本上，有没有办法可以单独创建和删除测试文件而不涉及测试本身？我尝试使用fixture，但它并没有真正起作用（也没有太多意义）。只是一个普通的函数可能吗？ - ShockDoctor

Answer 1

我认为问题在于您的测试输入和期望输出强烈地绑定在两个不同的位置，一个在参数中，另一个在测试代码中。
如果更改一个参数，则需要更改测试方法体，这样做是不正确的，而且还会产生重复的代码。
我认为您应该使用参数test(test_data, expected output)，并将输入注入临时文件中。
然后调用函数并比较预期输出和实际输出。

@pytest.mark.parametrize('test_data, expected', [
    ([['Data0', 'A0A0A0', 1, 'Data3', 'Data4', 'Data5', 4000]],
      pd.DataFrame([['A0A0A0', 1, 4000]], columns=['Postal_Code', 'Store_Num', 'Sales'])),
    ([['Data0', 'A0A0A0', 1, 'Data3', 'Data4', 'Data5', 'Data6', 4000]],
      pd.DataFrame([['A0A0A0', 1, 4000]], columns=['Postal_Code', 'Store_Num', 'Sales']))
])
def test_sales_dataframe(test_data, expected):

    # Write your test data in a temporary file
    tmp_file = r'Path\To\Project\Output\Folder\tmp.csv';
    with open(tmp_file, 'w') as file:
        writer = csv.writer(file)
        writer.writerows(test_data)

    # Process the data
    sales_df = mks_sales_dataframe(tmp_file)

    # Compare expected and actual output
    testing.assert_frame_equal(expected, sales_df)

    # Clean the temporary file
    os.remove(tmp_file)

您也可以创建您自己的 .csv 文件，并将它们添加为测试资源，但是您将会有不同的输入和期望输出位置，这并不那么理想。

- Arnaud Claudel

---

如果您不介意的话，能否用一个简短的例子来说明一下？抱歉，我不太明白这行代码是指什么：“我认为你应该有参数test(input, expected output)并将输入注入到临时文件中。” - ShockDoctor

@ShockDoctor 我已经更新了我的答案。个人而言，我会删除 @parameterize，因为它不太易读，并且阻止您轻松添加更长的测试用例。 - Arnaud Claudel

太棒了，这正是我想要的！一旦我看到代码，你之前写的东西就豁然开朗了。 - ShockDoctor

Answer 2

通过在 TestCase 中使用 setUp 和 tearDown 方法，可以减少一些重复。

import os
import csv
import unittest

test_file = 'test.csv'
rows = [
    ['0a', '0b', '0c'],
    ['1a', '1b', '1c'],
]


class TestCsv(unittest.TestCase):

    def setUp(self):
        with open(test_file, 'w', newline='') as csv_file:
            writer = csv.writer(csv_file, dialect='excel')
            writer.writerows(rows)

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

    def test_read_line(self):
        with open(test_file, 'r') as csv_file:
            reader = csv.reader(csv_file, dialect='excel')
            self.assertEqual(next(reader), rows[0])
            self.assertEqual(next(reader), rows[1])


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

- gman3g