diff --git a/src/SymptomHistogramByBatchcodeTableFactory.py b/src/SymptomHistogramByBatchcodeTableFactory.py index 53f3c023a84..8df21f09ffd 100644 --- a/src/SymptomHistogramByBatchcodeTableFactory.py +++ b/src/SymptomHistogramByBatchcodeTableFactory.py @@ -2,4 +2,7 @@ class SymptomHistogramByBatchcodeTableFactory: @staticmethod def createSymptomHistogramByBatchcodeTable(symptomByBatchcodeTable): - return symptomByBatchcodeTable.groupby(['VAX_LOT1'])['SYMPTOM'].value_counts().to_frame('SYMPTOM_COUNT_BY_VAX_LOT') + return (symptomByBatchcodeTable + .groupby(symptomByBatchcodeTable.index.names) + ['SYMPTOM'].value_counts() + .to_frame('SYMPTOM_COUNT_BY_VAX_LOT')) diff --git a/src/SymptomHistogramByBatchcodeTableFactoryTest.py b/src/SymptomHistogramByBatchcodeTableFactoryTest.py index 37429b852a6..8d518a69bcc 100644 --- a/src/SymptomHistogramByBatchcodeTableFactoryTest.py +++ b/src/SymptomHistogramByBatchcodeTableFactoryTest.py @@ -33,3 +33,29 @@ class SymptomHistogramByBatchcodeTableFactoryTest(unittest.TestCase): names = ['VAX_LOT1', 'SYMPTOM'], tuples = [['1808982', 'Blood pressure orthostatic abnormal'], ['EW0175', 'Blood pressure orthostatic abnormal']]))) + + def test_createSymptomHistogramByBatchcodeTable_two_VAX_LOTs_Index(self): + # Given + symptomByBatchcodeTable = TestHelper.createDataFrame( + columns = ['SYMPTOM'], + data = [ ['Blood pressure orthostatic abnormal'], + ['Blood pressure orthostatic abnormal'], + ['Headache']], + index = pd.MultiIndex.from_tuples( + names = ['VAX_LOT1', 'VAX_LOT2'], + tuples = [['1808982', 'EW0175']] * 3)) + + # When + symptomHistogramByBatchcodeTable = SymptomHistogramByBatchcodeTableFactory.createSymptomHistogramByBatchcodeTable(symptomByBatchcodeTable) + + # Then + assert_frame_equal( + symptomHistogramByBatchcodeTable, + TestHelper.createDataFrame( + columns = ['SYMPTOM_COUNT_BY_VAX_LOT'], + data = [ [2], + [1]], + index = pd.MultiIndex.from_tuples( + names = ['VAX_LOT1', 'VAX_LOT2', 'SYMPTOM'], + tuples = [['1808982', 'EW0175', 'Blood pressure orthostatic abnormal'], + ['1808982', 'EW0175', 'Headache']])))