CODEWITHSUNDEEP
pythonpandasunit-testingmocking
maslak
maslak

Reputation: 1255

How to test if saved csv file is correct

I have a small function which returns nothing but creates csv file instead. I would like to unittest it if given output csv file contains proper data. How can I achieve that?

function

def create_csv(input, output)
    df = pd.read_csv(input, names=['number1', 'number2'])
    df.number1 *= 2
    df.number2 *= 5
    df.to_csv(output, index=False, header=False)

Upvotes: 1

Views: 2181

Answers (2)

André C. Andersen
André C. Andersen

Reputation: 9405

Test the data frame just before you persist it to a CSV file. The NumPy test scripts and Pandas test functions can be helpful. If you feel you need to test the CSV file itself, you are actually distrusting pandas' to_csv() function. The best course of action then is to review the unit tests of pandas, and if you feel it is lacking do a pull request to amend it with the tests you feel are missing.

Upvotes: 2

wwii
wwii

Reputation: 23763

.to_csv and .read_csv take buffers as well as string paths. The docs for those methods suggest using io.StringIO as an alternative to an open file or file path.

In the unit test, use io.StringIO objects to pass to the function during the test. Something like this for a simple test of the function

import unittest, io
class TestOne(unittest.TestCase):
    def setUp(self):
        s = '''3,4\n2,1\n5,6'''
        self.fakecsv = io.StringIO(s)
        self.fakefile = io.StringIO()
        self.result = '''6,20\n4,5\n10,30\n'''

    def test_create(self):
        create_csv(self.fakecsv, self.fakefile)
        self.fakefile.seek(0)
        self.assertEqual(self.fakefile.read(), self.result)

Upvotes: 2

Related Questions