Update test for the new changes

cdplib/FlattenData.py

@@ -27,20 +27,23 @@ class FlattenData():
 
         assert(isinstance(data, (list, dict, pd.DataFrame))),\
             "Parameter 'data' either be of List, Dictionary or DataFrame type"
-
+        in_length=0
         start = time.time()
         if type(data) is pd.DataFrame:
+            in_length = len(data.columns)
             return_data = self.flatten_dataframe(data)
-        if type(data) is dict:
+        elif type(data) is dict:
+            in_length = len(data)
             return_data = self.flatten_dict(data)
-        if type(data) is list:
+        elif type(data) is list:
+            in_length = len(data)
             return_data =  self.flatten_list(data)
         else:
-            self._log.log_and_raise_warning('Input data type is not supported')
+            self._log.log_and_raise_warning(("Input data type '{}' is not supported").format(type(data)))
             return None
 
         result_dataframe = pd.DataFrame.from_dict(return_data, orient='index')
-        self._log.info(('Data has been flattened, created {} columns in {} seconds').format(len(result_dataframe.columns)- len(data.columns), time.time()-start))
+        self._log.info(('Data has been flattened, created {} columns in {} seconds').format(len(result_dataframe.columns)- in_length, time.time()-start))
         return result_dataframe
 
     def flatten_dataframe(self, dataframe: pd.DataFrame, incoming_key: str = None):
@@ -73,7 +76,7 @@ class FlattenData():
 
             result_dict[index] = copy.deepcopy(temp_result_dict)
 
-        return result_dataframe
+        return result_dict
 
     def flatten_dict(self, dictionary: dict, incoming_key: str = None):
         '''

cdplib/unit_tests/TestFlattenData.py

@@ -34,10 +34,10 @@ class TestMongodbHandler(unittest.TestCase):
 
         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'])
+        flattened_df = self.flattener.flatten(nested_data_df)        
+        
+        self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_dict.loc['two_levels_one_level', 0])
+        self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_list.loc['0_two_levels_one_level', 0])
         self.assertEqual(nested_data_dict["two_levels"]["one_level"], flattened_df.loc[0 , 'two_levels_one_level'])
         
 if __name__ == '__main__':