test_read_write_csv.py

import unittest
import pandas as pd
from ninety_nine_problems.one_to_twenty import OneToTwenty


class TestReadWriteCsv(unittest.TestCase):
    def test_read_write(self):
        # 1 Write data to file
        data = OneToTwenty.p10_number_of_duplicates(['a', 'a', 'a', 'a', 'c', 'c', 'c', 'c', 'b', 'b', 'b', 'b',
                                                     'b', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'd', 'e', 'e',
                                                     'e', 'e', 'f', 'f', 'f', 'f', 'f', 'f', 'f', 'f', 'g', 'h',
                                                     'i', 'j', 'j', 'j', 'j', 'j', 'k', 'k', 'k', 'k', 'k', 'k',
                                                     'k', 'k', 'k', 'k', 'k', 'k', 'k', 'k', 'k'])
        unread_data = pd.DataFrame(data, columns=['count', 'char'])
        print(unread_data)
        unread_data.to_csv('test_data.csv', index=False)

        # 2 Read data from file
        read_data = pd.read_csv('test_data.csv')

        # 3 Make sure data from 1 is equal to 2
        self.assertTrue(read_data.equals(unread_data))

    def test_first_row(self):
        # 1 Read the first row of file
        read_data = pd.read_csv('test_data.csv')
        first_row_count = read_data.iloc[0, 0]
        first_row_char = read_data.iloc[0, 1]
        # 2 Make sure 1st row is equal to [4 , 'a']
        self.assertEqual(first_row_count, 4)
        self.assertEqual(first_row_char, 'a')

    def test_last_row(self):
        # 1 Read the last row of file
        read_data = pd.read_csv('test_data.csv')
        last_row_count = read_data.iloc[-1, 0]
        last_row_char = read_data.iloc[-1, -1]
        # 2 Make sure last row is equal to [15, 'k']
        self.assertEqual(last_row_count, 15)
        self.assertEqual(last_row_char, 'k')


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