refining SymptomsByBatchcodesTableFactoryTest

src/SymptomsByBatchcodesTableFactory.py

@@ -5,12 +5,20 @@ class SymptomsByBatchcodesTableFactory:
 
     @staticmethod
     def createSymptomsByBatchcodesTable(VAERSVAX, VAERSSYMPTOMS):
-        index_columns = ['VAX_LOT1', 'VAX_LOT2']
+        index_columns = SymptomsByBatchcodesTableFactory._getIndexColumns(VAERSVAX)
         return pd.merge(
             SymptomsByBatchcodesTableFactory._get_VAERSVAX_WITH_VAX_LOTS(VAERSVAX, index_columns),
             SymptomsByBatchcodesTableFactory._getSymptomsTable(VAERSSYMPTOMS),
             on = 'VAERS_ID').set_index(index_columns)[['SYMPTOMS']]
     
+    @staticmethod
+    def _getIndexColumns(VAERSVAX):
+        return [f"VAX_LOT{num}" for num in range(1, SymptomsByBatchcodesTableFactory._getMaxNumShots(VAERSVAX) + 1)]
+
+    @staticmethod
+    def _getMaxNumShots(VAERSVAX):
+        return VAERSVAX.index.value_counts().iloc[0]
+
     @staticmethod
     def _get_VAERSVAX_WITH_VAX_LOTS(VAERSVAX, index_columns):
         return pd.concat(
@@ -21,7 +29,7 @@ class SymptomsByBatchcodesTableFactory:
     def _getVaxLotsTable(VAERSVAX, index_columns):
         VAX_LOT_LIST_Table = VAERSVAX.groupby("VAERS_ID").agg(VAX_LOT_LIST = pd.NamedAgg(column = 'VAX_LOT', aggfunc = list))
         return pd.DataFrame(
-            [fill(VAX_LOTS, 2, str(np.nan)) for VAX_LOTS in VAX_LOT_LIST_Table['VAX_LOT_LIST'].tolist()],
+            [fill(VAX_LOTS, len(index_columns), str(np.nan)) for VAX_LOTS in VAX_LOT_LIST_Table['VAX_LOT_LIST'].tolist()],
             columns = index_columns,
             index = VAX_LOT_LIST_Table.index)
 

src/SymptomsByBatchcodesTableFactoryTest.py

@@ -92,9 +92,10 @@ class SymptomsByBatchcodesTableFactoryTest(unittest.TestCase):
                 columns = ['SYMPTOMS'],
                 data = [  ['Blood pressure orthostatic abnormal'],
                           ['Blood pressure orthostatic abnormal']],
-                index = pd.MultiIndex.from_tuples(
-                    names =   ['VAX_LOT1', 'VAX_LOT2'],
-                    tuples = [['EW0175', str(np.nan)]] * 2)),
+                index = pd.Index(
+                    name = 'VAX_LOT1',
+                    data = ['EW0175',
+                            'EW0175'])),
                 check_dtype = False)
 
     def test_createSymptomsByBatchcodesTable_two_patients_distinct_symptoms(self):