adding test CountryCountsByClickedBatchcodeProviderTest.test_getCountryCountsByClickedBatchcode_fromCityResolution()

src/CountryCountsByClickedBatchcodeProvider.py

@@ -4,16 +4,31 @@ class CountryCountsByClickedBatchcodeProvider:
 
     @staticmethod
     def getCountryCountsByClickedBatchcode(file):
-        return CountryCountsByClickedBatchcodeProvider._read_csv(
+        if CountryCountsByClickedBatchcodeProvider._hasCityColumn(file):
-            file = file,
+            return CountryCountsByClickedBatchcodeProvider._getCountryCountsByClickedBatchcode_fromCityResolution(file)
-            columns = {
+        else:
-                'Country': 'COUNTRY',
+            return CountryCountsByClickedBatchcodeProvider._getCountryCountsByClickedBatchcode_fromCountryResolution(file)
-                'Event count': 'COUNTRY_COUNT_BY_VAX_LOT'
+
-            },
-            index_columns = ['COUNTRY'])
-    
     @staticmethod
-    def getCityCountsByClickedBatchcode(file):
+    def _hasCityColumn(file):
+        return 'City' in CountryCountsByClickedBatchcodeProvider._read_raw_csv(file).columns
+
+    @staticmethod
+    def _read_raw_csv(file):
+        return pd.read_csv(file, index_col = 0, skiprows = [0, 1, 2, 3, 4, 5, 7])
+
+    @staticmethod
+    def _getCountryCountsByClickedBatchcode_fromCityResolution(file):
+        cityCountsByClickedBatchcodeTable = CountryCountsByClickedBatchcodeProvider._getCityCountsByClickedBatchcode(file)
+        return (cityCountsByClickedBatchcodeTable
+                .groupby(['VAX_LOT', 'COUNTRY'])
+                .agg(COUNTRY_COUNT_BY_VAX_LOT =
+                      pd.NamedAgg(
+                          column = 'CITY_COUNT_BY_VAX_LOT',
+                          aggfunc = sum)))
+
+    @staticmethod
+    def _getCityCountsByClickedBatchcode(file):
         return CountryCountsByClickedBatchcodeProvider._read_csv(
             file = file,
             columns = {
@@ -23,13 +38,24 @@ class CountryCountsByClickedBatchcodeProvider:
                 'Event count': 'CITY_COUNT_BY_VAX_LOT'
             },
             index_columns = ['COUNTRY', 'REGION', 'CITY'])
+
+    @staticmethod
+    def _getCountryCountsByClickedBatchcode_fromCountryResolution(file):
+        return CountryCountsByClickedBatchcodeProvider._read_csv(
+            file = file,
+            columns = {
+                'Country': 'COUNTRY',
+                'Event count': 'COUNTRY_COUNT_BY_VAX_LOT'
+            },
+            index_columns = ['COUNTRY'])
     
     @staticmethod
     def _read_csv(file, columns, index_columns):
-        exploration = pd.read_csv(file, index_col = 0, skiprows = [0, 1, 2, 3, 4, 5, 7])
+        dataframe = CountryCountsByClickedBatchcodeProvider._read_raw_csv(file)
-        exploration.index.name = 'VAX_LOT'
+        dataframe.index.name = 'VAX_LOT'
-        exploration.rename(
+        dataframe.rename(
             columns = columns,
             inplace = True)
-        exploration.set_index(index_columns, append = True, inplace = True)
+        dataframe.set_index(index_columns, append = True, inplace = True)
-        return exploration
+        return dataframe
+    

src/CountryCountsByClickedBatchcodeProviderTest.py

@@ -6,9 +6,9 @@ from CountryCountsByClickedBatchcodeProvider import CountryCountsByClickedBatchc
 
 class CountryCountsByClickedBatchcodeProviderTest(unittest.TestCase):
 
-    def test_getCountryCountsByClickedBatchcode(self):
+    def test_getCountryCountsByClickedBatchcode_fromCountryResolution(self):
         # Given
-            
+
         # When
         countryCountsByClickedBatchcodeTable = CountryCountsByClickedBatchcodeProvider.getCountryCountsByClickedBatchcode('src/testdata/GoogleAnalytics/CountryByBatchcode 20230302-20230430.csv')
         
@@ -28,9 +28,9 @@ class CountryCountsByClickedBatchcodeProviderTest(unittest.TestCase):
 
     def test_getCityCountsByClickedBatchcode(self):
         # Given
-            
+
         # When
-        cityCountsByClickedBatchcodeTable = CountryCountsByClickedBatchcodeProvider.getCityCountsByClickedBatchcode('src/testdata/GoogleAnalytics/CountryByBatchcode 20230730-20230929.csv')
+        cityCountsByClickedBatchcodeTable = CountryCountsByClickedBatchcodeProvider._getCityCountsByClickedBatchcode('src/testdata/GoogleAnalytics/CountryByBatchcode 20230730-20230929.csv')
         
         # Then
         assert_frame_equal(
@@ -45,3 +45,21 @@ class CountryCountsByClickedBatchcodeProviderTest(unittest.TestCase):
                             tuples = [['#003B21A', 'United States', 'California', 'Roseville'],
                                       ['000086A',  'Germany',       'Bavaria',    'Nordlingen'],
                                       ['000086A',  'Germany',       'Bavaria',    'Nuremberg']])))
+
+    def test_getCountryCountsByClickedBatchcode_fromCityResolution(self):
+        # Given
+
+        # When
+        countryCountsByClickedBatchcodeTable = CountryCountsByClickedBatchcodeProvider.getCountryCountsByClickedBatchcode('src/testdata/GoogleAnalytics/CountryByBatchcode 20230730-20230929.csv')
+        
+        # Then
+        assert_frame_equal(
+            countryCountsByClickedBatchcodeTable,
+            TestHelper.createDataFrame(
+                columns = ['COUNTRY_COUNT_BY_VAX_LOT'],
+                data = [  [100],
+                          [10 + 20]],
+                index = pd.MultiIndex.from_tuples(
+                            names =   ['VAX_LOT',  'COUNTRY'],
+                            tuples = [['#003B21A', 'United States'],
+                                      ['000086A',  'Germany']])))