diff --git a/src/TableByBatchcodeFilter.py b/src/TableByBatchcodeFilter.py new file mode 100644 index 00000000000..37ae2511d33 --- /dev/null +++ b/src/TableByBatchcodeFilter.py @@ -0,0 +1,9 @@ +class TableByBatchcodeFilter: + + @staticmethod + def filterTableByBatchcode(batchcode, table): + table = table.reset_index() + filteredTable = table[ + (table['VAX_LOT1'] == batchcode) | + (table['VAX_LOT2'] == batchcode)] + return filteredTable.set_index(['VAX_LOT1', 'VAX_LOT2']) diff --git a/src/TableByBatchcodeFilterTest.py b/src/TableByBatchcodeFilterTest.py new file mode 100644 index 00000000000..0df8a3d51cd --- /dev/null +++ b/src/TableByBatchcodeFilterTest.py @@ -0,0 +1,37 @@ +import unittest +from pandas.testing import assert_frame_equal +from TableByBatchcodeFilter import TableByBatchcodeFilter +from TestHelper import TestHelper +import pandas as pd + +class TableByBatchcodeFilterTest(unittest.TestCase): + + def test_convertHistogramTable2JsonTable(self): + # Given + batchcode = '1808982' + symptomHistogramByBatchcodeTable = TestHelper.createDataFrame( + columns = ['SYMPTOM_COUNT_BY_VAX_LOT'], + data = [ ['{"Blood pressure orthostatic abnormal":5,"Chest discomfort":1}'], + ['{"Chest discomfort":2}'], + ['{"Chills":5}']], + index = pd.MultiIndex.from_tuples( + names = ['VAX_LOT1', 'VAX_LOT2'], + tuples = [[batchcode, 'EW0175'], + ['015M20A', batchcode], + ['015M20A', 'EW0175']])) + + # When + filteredTable = TableByBatchcodeFilter.filterTableByBatchcode(batchcode, symptomHistogramByBatchcodeTable) + + # Then + assert_frame_equal( + filteredTable, + TestHelper.createDataFrame( + columns = ['SYMPTOM_COUNT_BY_VAX_LOT'], + data = [ ['{"Blood pressure orthostatic abnormal":5,"Chest discomfort":1}'], + ['{"Chest discomfort":2}']], + index = pd.MultiIndex.from_tuples( + names = ['VAX_LOT1', 'VAX_LOT2'], + tuples = [[batchcode, 'EW0175'], + ['015M20A', batchcode]]))) +