adding test CountryCountsByClickedBatchcodeProviderTest.test_getCountryCountsByClickedBatchcode_fromCityResolution()

src/CountryCountsByClickedBatchcodeProvider.py

@@ -4,16 +4,31 @@ class CountryCountsByClickedBatchcodeProvider:
 
     @staticmethod
     def getCountryCountsByClickedBatchcode(file):
-        return CountryCountsByClickedBatchcodeProvider._read_csv(
-            file = file,
-            columns = {
-                'Country': 'COUNTRY',
-                'Event count': 'COUNTRY_COUNT_BY_VAX_LOT'
-            },
-            index_columns = ['COUNTRY'])
-    
+        if CountryCountsByClickedBatchcodeProvider._hasCityColumn(file):
+            return CountryCountsByClickedBatchcodeProvider._getCountryCountsByClickedBatchcode_fromCityResolution(file)
+        else:
+            return CountryCountsByClickedBatchcodeProvider._getCountryCountsByClickedBatchcode_fromCountryResolution(file)
+
     @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])
-        exploration.index.name = 'VAX_LOT'
-        exploration.rename(
+        dataframe = CountryCountsByClickedBatchcodeProvider._read_raw_csv(file)
+        dataframe.index.name = 'VAX_LOT'
+        dataframe.rename(
             columns = columns,
             inplace = True)
-        exploration.set_index(index_columns, append = True, inplace = True)
-        return exploration
+        dataframe.set_index(index_columns, append = True, inplace = True)
+        return dataframe
+