filterByBatchcodes()

src/CountriesByBatchcodeProviderTest.py

@@ -0,0 +1,38 @@
+import unittest
+from TestHelper import TestHelper
+from pandas.testing import assert_frame_equal
+import pandas as pd
+from CountriesByBatchcodeProvider import filterByBatchcodes
+
+class CountriesByBatchcodeProviderTest(unittest.TestCase):
+    
+    def test_filterByBatchcodes(self):
+        # Given
+        countryCountsByBatchcode = TestHelper.createDataFrame(
+                columns = ['COUNTRY_COUNT_BY_VAX_LOT Clicked', 'COUNTRY_COUNT_BY_VAX_LOT Before Deletion'],
+                data = [  [10,                                 20],
+                          [90,                                 95],
+                          [15,                                 30],
+                          [70,                                 80]],
+                index = pd.MultiIndex.from_tuples(
+                            names =   ['VAX_LOT',   'COUNTRY'],
+                            tuples = [('!D0181',    'Germany'),
+                                      ('# 009C01A', 'United States'),
+                                      ('!D0181',    'Hungary'),
+                                      ('# 009C01A', 'Germany')]))
+
+        # When
+        countryCountsByBatchcodeFiltered = filterByBatchcodes(countryCountsByBatchcode, ['!D0181'])
+
+        # Then
+        assert_frame_equal(
+            countryCountsByBatchcodeFiltered,
+            TestHelper.createDataFrame(
+                columns = ['COUNTRY_COUNT_BY_VAX_LOT Clicked', 'COUNTRY_COUNT_BY_VAX_LOT Before Deletion'],
+                data = [  [10,                                 20],
+                          [15,                                 30]],
+                index = pd.MultiIndex.from_tuples(
+                            names =   ['VAX_LOT',   'COUNTRY'],
+                            tuples = [('!D0181',    'Germany'),
+                                      ('!D0181',    'Hungary')])))
+