refining TableByBatchcodeFilterTest

src/TableByBatchcodeFilter.py

@@ -1,9 +1,18 @@
+from functools import reduce
+
+
 class TableByBatchcodeFilter:
 
     @staticmethod
     def filterTableByBatchcode(batchcode, table):
+        batchcodeColumns = table.index.names
         table = table.reset_index()
-        filteredTable = table[
-            (table['VAX_LOT1'] == batchcode) |
-            (table['VAX_LOT2'] == batchcode)]
-        return filteredTable.set_index(['VAX_LOT1', 'VAX_LOT2'])
+        filteredTable = table[TableByBatchcodeFilter._existsBatchcodeInAnyBatchcodeColumn(table, batchcodeColumns, batchcode)]
+        return filteredTable.set_index(batchcodeColumns)
+
+    @staticmethod
+    def _existsBatchcodeInAnyBatchcodeColumn(table, batchcodeColumns, batchcode):
+        return reduce(
+            lambda accum, batchcodeColumn: accum | (table[batchcodeColumn] == batchcode),
+            batchcodeColumns,
+            [False] * len(table.index))