adding test SymptomHistogramByBatchcodeTableFactoryTest.test_createSymptomHistogramByBatchcodeTable_two_VAX_LOTs_Index()

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'))

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']])))