cdplib/cdplib/unit_tests/TestFlattenData.py

import unittest
import sys
import os
import pandas as pd
sys.path.append(os.getcwd())
from cdplib.log import Log
from cdplib.FlattenData import FlattenData

class TestMongodbHandler(unittest.TestCase):

    def setUp(self):
        self.flattener = FlattenData()

    def test_A_flatten(self):
        '''
        Create some nested test data, in the formats: dict, list and dataframe
        Flatten the test data
        Compare the results
        '''
        nested_data_dict = {
            "one_level": "test_level_1",
            "two_levels": {
                "one_level": "test_level_2"
            },
            "three_levels": {
                "two_levels": {
                    "one_level": "test_level_3"
                }
            }
        }

        nested_data_list = [nested_data_dict, nested_data_dict]
        nested_data_df = pd.DataFrame.from_dict([nested_data_dict])

        flattened_dict = self.flattener.flatten(nested_data_dict)
        flattened_list = self.flattener.flatten(nested_data_list)
        flattened_df = self.flattener.flatten(nested_data_df)

        self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_dict['two_levels_one_level'])
        self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_list['0_two_levels_one_level'])
        self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_df.loc[0 , 'two_levels_one_level'])
        
if __name__ == '__main__':
    unittest.main()