-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_import_cleaning.py
More file actions
38 lines (30 loc) · 1.5 KB
/
test_import_cleaning.py
File metadata and controls
38 lines (30 loc) · 1.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
### Testing of import_cleaning function
import unittest
from import_cleaning import *
import pickle
data = pd.read_pickle('dataframe.pkl')
columns =data.columns
# unittest test case
class StringConverting(unittest.TestCase):
def test_convert_column_types(self):
# Creating a sample DataFrame for testing
sample_data = pd.DataFrame({
'sample_col1': ['NaN', 'string1', 'string2'],
'sample_col2': ['string3', 'NaN', 'string4']
})
# Call the function to be tested
result = convert_column_types(sample_data, ['sample_col1', 'sample_col2'])
# Check if 'Not Applicable' is replaced with pd.NA and column type is 'string'
#if we remove a letter in the expected output eg type "sting1", we get an error
self.assertEqual(result['sample_col1'].tolist(), [pd.NA, 'string1', 'string2'])
self.assertEqual(result['sample_col2'].tolist(), ['string3', pd.NA, 'string4'])
self.assertEqual(result['sample_col1'].dtype, 'string')
self.assertEqual(result['sample_col2'].dtype, 'string')
# Check if the number of columns in our actual dataframe matches our desired output
assert (len(columns) == 16)
# Check that no more 'NaN' nor 'Not Applicable' are in the actual dataframe
for col in data:
self.assertFalse(((data == 'Not Applicable').any().any())
&(data == 'NaN').any().any())
if __name__ =='__main__':
unittest.main()