Merge pull request #1108 from cmu-delphi/acquisition_cronicle_change

Acquisition change for run in cronicle

Author: minhkhul

integrations/acquisition/covidcast/test_csv_uploading.py

@@ -4,7 +4,7 @@
 from datetime import date
 import os
 import unittest
-from unittest.mock import MagicMock
+import argparse
 
 # third party
 import mysql.connector
@@ -14,7 +14,7 @@
 # first party
 from delphi_utils import Nans
 from delphi.epidata.client.delphi_epidata import Epidata
-from delphi.epidata.acquisition.covidcast.csv_to_database import main
+from delphi.epidata.acquisition.covidcast.csv_to_database import main, get_argument_parser
 import delphi.operations.secrets as secrets
 
 # py3tester coverage target (equivalent to `import *`)
@@ -92,16 +92,12 @@ def test_uploading(self):
 
     # make some fake data files
     data_dir = 'covid/data'
-    source_receiving_dir = data_dir + '/receiving/src-name'
+    indicator_dir = 'src-name'
+    source_receiving_dir = data_dir + '/receiving/' + indicator_dir
     log_file_directory = "/var/log/"
     os.makedirs(source_receiving_dir, exist_ok=True)
     os.makedirs(log_file_directory, exist_ok=True)
-    # TODO: use an actual argparse object for the args instead of a MagicMock
-    args = MagicMock(
-          log_file=log_file_directory +
-          "output.log",
-          data_dir=data_dir,
-          specific_issue_date=False)
+    args = get_argument_parser().parse_args(["--log_file", log_file_directory + "output.log", "--data_dir", data_dir, "--indicator_name", indicator_dir])
     uploader_column_rename = {"geo_id": "geo_value", "val": "value", "se": "stderr", "missing_val": "missing_value", "missing_se": "missing_stderr"}
 
 

src/acquisition/covidcast/csv_importer.py

@@ -133,10 +133,11 @@ def find_issue_specific_csv_files(scan_dir):
 
 
   @staticmethod
-  def find_csv_files(scan_dir, issue=(date.today(), epi.Week.fromdate(date.today()))):
+  def find_csv_files(scan_dir, issue=(date.today(), epi.Week.fromdate(date.today())), indicator_dir= "*"):
     """Recursively search for and yield covidcast-format CSV files.
 
     scan_dir: the directory to scan (recursively)
+    indicator_dir: specify one indicator with .csv files inside
 
     The return value is a tuple of (path, details), where, if the path was
     valid, details is a tuple of (source, signal, time_type, geo_type,
@@ -149,7 +150,7 @@ def find_csv_files(scan_dir, issue=(date.today(), epi.Week.fromdate(date.today()
     issue_value=-1
     lag_value=-1
 
-    for path in sorted(glob(os.path.join(scan_dir, '*', '*'))):
+    for path in sorted(glob(os.path.join(scan_dir, indicator_dir, '*'))):
       # safe to ignore this file
       if not path.lower().endswith('.csv'):
         continue

src/acquisition/covidcast/csv_to_database.py

@@ -28,16 +28,21 @@ def get_argument_parser():
   parser.add_argument(
     '--log_file',
     help="filename for log output (defaults to stdout)")
+  parser.add_argument(
+    '--indicator_name',
+    nargs='?',
+    default='*',
+    help='Name of one indicator directory to run acquisition on')
   return parser
 
 
-def collect_files(data_dir: str, specific_issue_date: bool):
+def collect_files(data_dir: str, specific_issue_date: bool, indicator_name="*"):
   """Fetch path and data profile details for each file to upload."""
   logger= get_structured_logger('collect_files')
   if specific_issue_date:
     results = list(CsvImporter.find_issue_specific_csv_files(data_dir))
   else:
-    results = list(CsvImporter.find_csv_files(os.path.join(data_dir, 'receiving')))
+    results = list(CsvImporter.find_csv_files(os.path.join(data_dir, 'receiving'), indicator_dir=indicator_name))
   logger.info(f'found {len(results)} files')
   return results
 
@@ -146,9 +151,8 @@ def main(args):
 
   logger = get_structured_logger("csv_ingestion", filename=args.log_file)
   start_time = time.time()
-
   # shortcut escape without hitting db if nothing to do
-  path_details = collect_files(args.data_dir, args.specific_issue_date)
+  path_details=collect_files(args.data_dir, args.specific_issue_date, indicator_name = args.indicator_name)
   if not path_details:
     logger.info('nothing to do; exiting...')
     return