move pull date creation into pull_gs_data

Author: nmdefries

google_symptoms/delphi_google_symptoms/pull.py

@@ -9,6 +9,7 @@
 from google.oauth2 import service_account
 
 from .constants import COMBINED_METRIC, DC_FIPS, DTYPE_CONVERSIONS, METRICS, SYMPTOM_SETS
+from .date_utils import generate_query_dates
 
 # Create map of BigQuery symptom column names to desired column names.
 colname_map = {"symptom_" +
@@ -214,7 +215,7 @@ def initialize_credentials(credentials):
     pandas_gbq.context.project = credentials.project_id
 
 
-def pull_gs_data(credentials, export_date_range):
+def pull_gs_data(credentials, export_start_date, export_end_date, num_export_days, custom_run_flag):
     """Pull latest dataset for each geo level and combine.
 
     PS:  No information for PR
@@ -237,7 +238,7 @@ def pull_gs_data(credentials, export_date_range):
     dict: {"county": pd.DataFrame, "state": pd.DataFrame}
     """
     # Fetch and format dates we want to attempt to retrieve
-
+    export_date_range = generate_query_dates(export_start_date, export_end_date, num_export_days, custom_run_flag)
     retrieve_dates = format_dates_for_query(export_date_range)
 
     initialize_credentials(credentials)

google_symptoms/delphi_google_symptoms/run.py

@@ -12,7 +12,7 @@
 from delphi_utils import create_export_csv, get_structured_logger
 
 from .constants import COMBINED_METRIC, GEO_RESOLUTIONS, SMOOTHERS, SMOOTHERS_MAP
-from .date_utils import generate_num_export_days, generate_query_dates
+from .date_utils import generate_num_export_days
 from .geo import geo_map
 from .pull import pull_gs_data
 
@@ -58,10 +58,9 @@ def run_module(params, logger=None):
     custom_run_flag = (
         False if not params["indicator"].get("custom_run", False) else params["indicator"].get("custom_run", False)
     )
-    export_date_range = generate_query_dates(export_start_date, export_end_date, num_export_days, custom_run_flag)
 
     # Pull GS data
-    dfs = pull_gs_data(params["indicator"]["bigquery_credentials"], export_date_range)
+    dfs = pull_gs_data(params["indicator"]["bigquery_credentials"], export_start_date, export_end_date, num_export_days, custom_run_flag)
     for geo_res in GEO_RESOLUTIONS:
         if geo_res == "state":
             df_pull = dfs["state"]

google_symptoms/tests/test_pull.py

@@ -44,7 +44,8 @@ def test_good_file(self, mock_credentials, mock_read_gbq):
             "20201230", "%Y%m%d")
         end_date = datetime.combine(date.today(), datetime.min.time())
 
-        dfs = pull_gs_data("", [start_date, end_date])
+        dfs = pull_gs_data("", datetime.strptime(
+            "20201230", "%Y%m%d"), datetime.combine(date.today(), datetime.min.time()), 0, False)
 
         for level in ["county", "state"]:
             df = dfs[level]