Files
HowBadIsMyBatch/src/DataFrameFilterTest.py
frankknoll e87fe0c8ba refactoring
2022-11-22 12:40:26 +01:00

65 lines
3.0 KiB
Python

import unittest
from pandas.testing import assert_frame_equal
from VaersDescr2DataFrameConverter import VaersDescr2DataFrameConverter
from TestHelper import TestHelper
from DataFrameFilter import DataFrameFilter
class DataFrameFilterTest(unittest.TestCase):
def test_filterByCovid19(self):
# Given
dataFrame = VaersDescr2DataFrameConverter.createDataFrameFromDescrs(
[
{
'VAERSDATA': TestHelper.createDataFrame(
columns = ['DIED', 'L_THREAT', 'DISABLE'],
data = [ [1, 0, 0],
[0, 0, 1]],
index = [
"0916600",
"0916601"]),
'VAERSVAX': TestHelper.createDataFrame(
columns = ['VAX_TYPE', 'VAX_MANU', 'VAX_LOT', 'VAX_DOSE_SERIES'],
data = [ ['COVID19', 'MODERNA', '037K20A', '1'],
['COVID19', 'MODERNA', '025L20A', '1']],
index = [
"0916600",
"0916601"],
dtypes = {'VAX_DOSE_SERIES': "string"})
},
{
'VAERSDATA': TestHelper.createDataFrame(
columns = ['DIED', 'L_THREAT', 'DISABLE'],
data = [ [0, 0, 0],
[0, 0, 1]],
index = [
"1996873",
"1996874"]),
'VAERSVAX': TestHelper.createDataFrame(
columns = ['VAX_TYPE', 'VAX_MANU', 'VAX_LOT', 'VAX_DOSE_SERIES'],
data = [ ['HPV9', 'MERCK & CO. INC.', 'R017624', 'UNK'],
['COVID19', 'MODERNA', '025L20A', '1']],
index = [
"1996873",
"1996874"],
dtypes = {'VAX_DOSE_SERIES': "string"})
}
])
dataFrameFilter = DataFrameFilter()
# When
dataFrame = dataFrameFilter.filterByCovid19(dataFrame)
# Then
dataFrameExpected = TestHelper.createDataFrame(
columns = ['DIED', 'L_THREAT', 'DISABLE', 'VAX_TYPE', 'VAX_MANU', 'VAX_LOT', 'VAX_DOSE_SERIES'],
data = [ [1, 0, 0, 'COVID19', 'MODERNA', '037K20A', '1'],
[0, 0, 1, 'COVID19', 'MODERNA', '025L20A', '1'],
[0, 0, 1, 'COVID19', 'MODERNA', '025L20A', '1']],
index = [
"0916600",
"0916601",
"1996874"],
dtypes = {'VAX_DOSE_SERIES': "string"})
assert_frame_equal(dataFrame, dataFrameExpected, check_dtype = False)