suggested changes and refactor pull_nhsn_data

Author: aysim319

nhsn/delphi_nhsn/pull.py

@@ -123,7 +123,9 @@ def pull_nhsn_data(
     backup_dir: str,
     custom_run: bool,
     issue_date: Optional[str],
+    preliminary: bool = False,
     logger: Optional[logging.Logger] = None,
+
 ):
     """Pull the latest NHSN hospital admission data, and conforms it into a dataset.
 
@@ -140,6 +142,10 @@ def pull_nhsn_data(
         Directory to which to save raw backup data
     custom_run: bool
         Flag indicating if the current run is a patch. If so, don't save any data to disk
+    preliminary: bool
+        Flag indicating if the grabbing main or preliminary data
+    issue_date:
+        date to indicate which backup file to pull for patching
     logger: Optional[logging.Logger]
         logger object
 
@@ -148,22 +154,26 @@ def pull_nhsn_data(
     pd.DataFrame
         Dataframe as described above.
     """
+    dataset_id = PRELIM_DATASET_ID if preliminary else MAIN_DATASET_ID
     # Pull data from Socrata API
     df = (
-        pull_data(socrata_token, MAIN_DATASET_ID, backup_dir, logger)
+        pull_data(socrata_token, dataset_id, backup_dir, logger)
         if not custom_run
-        else pull_data_from_file(backup_dir, issue_date, logger, prelim_flag=False)
+        else pull_data_from_file(backup_dir, issue_date, logger, prelim_flag=preliminary)
     )
 
-    recently_updated = True if custom_run else check_last_updated(socrata_token, MAIN_DATASET_ID, logger)
+    recently_updated = True if custom_run else check_last_updated(socrata_token, dataset_id, logger)
+
+    type_dict = PRELIM_TYPE_DICT if preliminary else TYPE_DICT
+    keep_columns = list(type_dict.keys())
+    filtered_type_dict = copy.deepcopy(type_dict)
 
-    keep_columns = list(TYPE_DICT.keys())
+    signal_map = PRELIM_SIGNALS_MAP if preliminary else SIGNALS_MAP
 
     if not df.empty and recently_updated:
         df = df.rename(columns={"weekendingdate": "timestamp", "jurisdiction": "geo_id"})
-        filtered_type_dict = copy.deepcopy(TYPE_DICT)
 
-        for signal, col_name in SIGNALS_MAP.items():
+        for signal, col_name in signal_map.items():
             # older backups don't have certain columns
             try:
                 df[signal] = df[col_name]
@@ -181,66 +191,3 @@ def pull_nhsn_data(
         df = pd.DataFrame(columns=keep_columns)
 
     return df
-
-
-def pull_preliminary_nhsn_data(
-    socrata_token: str,
-    backup_dir: str,
-    custom_run: bool,
-    issue_date: Optional[str],
-    logger: Optional[logging.Logger] = None,
-):
-    """Pull the latest preliminary NHSN hospital admission data, and conforms it into a dataset.
-
-    The output dataset has:
-
-    - Each row corresponds to a single observation
-    - Each row additionally has columns for the signals in SIGNALS
-
-    Parameters
-    ----------
-    socrata_token: str
-        My App Token for pulling the NHSN data
-    backup_dir: str
-        Directory to which to save raw backup data
-    custom_run: bool
-        Flag indicating if the current run is a patch. If so, don't save any data to disk
-    logger: Optional[logging.Logger]
-        logger object
-
-    Returns
-    -------
-    pd.DataFrame
-        Dataframe as described above.
-    """
-    # Pull data from Socrata API
-    df = (
-        pull_data(socrata_token, PRELIM_DATASET_ID, backup_dir, logger)
-        if not custom_run
-        else pull_data_from_file(backup_dir, issue_date, logger, prelim_flag=True)
-    )
-
-    keep_columns = list(PRELIM_TYPE_DICT.keys())
-    recently_updated = True if custom_run else check_last_updated(socrata_token, PRELIM_DATASET_ID, logger)
-
-    if not df.empty and recently_updated:
-        df = df.rename(columns={"weekendingdate": "timestamp", "jurisdiction": "geo_id"})
-        filtered_type_dict = copy.deepcopy(PRELIM_TYPE_DICT)
-
-        for signal, col_name in PRELIM_SIGNALS_MAP.items():
-            try:
-                df[signal] = df[col_name]
-            except KeyError:
-                logger.info("column not available in data", col_name=col_name, signal=signal)
-                keep_columns.remove(signal)
-                del filtered_type_dict[signal]
-
-        df = df[keep_columns]
-        df = df.astype(filtered_type_dict)
-
-        df["geo_id"] = df["geo_id"].str.lower()
-        df.loc[df["geo_id"] == "usa", "geo_id"] = "us"
-    else:
-        df = pd.DataFrame(columns=keep_columns)
-
-    return df

nhsn/delphi_nhsn/run.py

@@ -24,7 +24,7 @@
 from delphi_utils.export import create_export_csv
 
 from .constants import GEOS, PRELIM_SIGNALS_MAP, SIGNALS_MAP
-from .pull import pull_nhsn_data, pull_preliminary_nhsn_data
+from .pull import pull_nhsn_data
 
 
 def run_module(params, logger=None):
@@ -56,8 +56,8 @@ def run_module(params, logger=None):
         export_start_date = export_start_date.strftime("%Y-%m-%d")
 
     nhsn_df = pull_nhsn_data(socrata_token, backup_dir, custom_run=custom_run, issue_date=issue_date, logger=logger)
-    preliminary_nhsn_df = pull_preliminary_nhsn_data(
-        socrata_token, backup_dir, custom_run=custom_run, issue_date=issue_date, logger=logger
+    preliminary_nhsn_df = pull_nhsn_data(
+        socrata_token, backup_dir, custom_run=custom_run, issue_date=issue_date, logger=logger, preliminary=True
     )
 
     geo_mapper = GeoMapper()

nhsn/tests/test_data/expected_df.csv

@@ -0,0 +1,21 @@
+timestamp,geo_id,confirmed_admissions_covid_ew,confirmed_admissions_flu_ew,confirmed_admissions_rsv_ew,hosprep_confirmed_admissions_covid_ew,hosprep_confirmed_admissions_flu_ew,hosprep_confirmed_admissions_rsv_ew
+2021-08-21,md,53.0,2.0,0.0,13.0,13.0,1.0
+2021-08-21,co,852.0,0.0,,92.0,78.0,0.0
+2021-08-21,us,10384.0,6049.0,84.0,5426.0,5426.0,469.0
+2021-08-28,co,835.0,1.0,,92.0,78.0,0.0
+2021-08-28,us,94596.0,262.0,,5391.0,4397.0,0.0
+2021-09-04,co,1000.0,3.0,,92.0,78.0,0.0
+2021-09-04,us,93241.0,282.0,,5392.0,4396.0,0.0
+2021-09-11,co,982.0,2.0,,92.0,78.0,0.0
+2021-09-11,us,88162.0,247.0,,5391.0,4377.0,0.0
+2021-09-18,co,955.0,0.0,,92.0,78.0,0.0
+2021-09-18,us,79169.0,261.0,,5394.0,4362.0,0.0
+2021-09-25,co,993.0,0.0,,92.0,78.0,0.0
+2021-09-25,us,67740.0,234.0,,5393.0,4368.0,0.0
+2021-10-02,co,970.0,0.0,,92.0,78.0,0.0
+2021-10-02,us,58076.0,253.0,,5395.0,4391.0,0.0
+2021-10-09,co,1079.0,1.0,,92.0,78.0,0.0
+2021-10-09,us,51744.0,341.0,,5396.0,4379.0,0.0
+2021-10-16,co,1231.0,0.0,,92.0,78.0,0.0
+2021-10-16,us,45978.0,266.0,,5394.0,4307.0,0.0
+2021-10-16,region 1,45978.0,266.0,,5394.0,4307.0,0.0

nhsn/tests/test_data/expected_df_prelim.csv

@@ -0,0 +1,20 @@
+timestamp,geo_id,confirmed_admissions_covid_ew_prelim,confirmed_admissions_flu_ew_prelim,confirmed_admissions_rsv_ew_prelim,hosprep_confirmed_admissions_covid_ew_prelim,hosprep_confirmed_admissions_flu_ew_prelim,hosprep_confirmed_admissions_rsv_ew_prelim
+2021-08-21,mi,269.0,523.0,1.0,152.0,152.0,4.0
+2021-08-21,co,852.0,0.0,,92.0,78.0,0.0
+2021-08-21,us,8946.0,5576.0,61.0,5422.0,5422.0,485.0
+2021-08-28,co,835.0,1.0,,92.0,78.0,0.0
+2021-08-28,us,94596.0,262.0,,5391.0,4397.0,0.0
+2021-09-04,co,1000.0,3.0,,92.0,78.0,0.0
+2021-09-04,us,93241.0,282.0,,5392.0,4396.0,0.0
+2021-09-11,co,982.0,2.0,,92.0,78.0,0.0
+2021-09-11,us,88162.0,247.0,,5391.0,4377.0,0.0
+2021-09-18,co,955.0,0.0,,92.0,78.0,0.0
+2021-09-18,us,79169.0,261.0,,5394.0,4362.0,0.0
+2021-09-25,co,993.0,0.0,,92.0,78.0,0.0
+2021-09-25,us,67740.0,234.0,,5393.0,4368.0,0.0
+2021-10-02,co,970.0,0.0,,92.0,78.0,0.0
+2021-10-02,us,58076.0,253.0,,5395.0,4391.0,0.0
+2021-10-09,co,1079.0,1.0,,92.0,78.0,0.0
+2021-10-09,us,51744.0,341.0,,5396.0,4379.0,0.0
+2021-10-16,co,1231.0,0.0,,92.0,78.0,0.0
+2021-10-16,us,45978.0,266.0,,5394.0,4307.0,0.0

nhsn/tests/test_patch.py

@@ -13,7 +13,7 @@
 from delphi_nhsn.patch import filter_source_files, patch
 from delphi_nhsn.constants import TOTAL_ADMISSION_COVID_COL, TOTAL_ADMISSION_FLU_COL, \
     NUM_HOSP_REPORTING_FLU_COL, NUM_HOSP_REPORTING_COVID_COL, GEOS, TOTAL_ADMISSION_COVID, TOTAL_ADMISSION_FLU, \
-    NUM_HOSP_REPORTING_COVID, NUM_HOSP_REPORTING_FLU, NUM_HOSP_REPORTING_RSV_COL, TOTAL_ADMISSION_RSV_COL
+     NUM_HOSP_REPORTING_RSV_COL, TOTAL_ADMISSION_RSV_COL
 from conftest import TEST_DATA, PRELIM_TEST_DATA, TEST_DIR
 
 class TestPatch:

nhsn/tests/test_pull.py

@@ -10,7 +10,7 @@
     pull_nhsn_data,
     pull_data,
     pull_data_from_file,
-    pull_preliminary_nhsn_data, check_last_updated
+    check_last_updated
 )
 from delphi_nhsn.constants import TYPE_DICT, PRELIM_TYPE_DICT, PRELIM_DATASET_ID, MAIN_DATASET_ID
 
@@ -21,12 +21,16 @@
              "test_data": TEST_DATA,
              "msg_prefix": "",
              "prelim_flag": False,
+             "expected_data": f"{TEST_DIR}/test_data/expected_df.csv",
+             "type_dict": TYPE_DICT,
              },
 
             {"id":PRELIM_DATASET_ID,
              "test_data":PRELIM_TEST_DATA,
              "msg_prefix": "Preliminary ",
              "prelim_flag": True,
+             "expected_data": f"{TEST_DIR}/test_data/expected_df_prelim.csv",
+             "type_dict": PRELIM_TYPE_DICT,
              }
             ]
 
@@ -75,123 +79,63 @@ def test_pull_from_file(self, caplog, dataset, params_w_patch):
 
     @patch("delphi_nhsn.pull.Socrata")
     @patch("delphi_nhsn.pull.create_backup_csv")
-    def test_pull_nhsn_data_output(self, mock_create_backup, mock_socrata, caplog, params):
+    @pytest.mark.parametrize('dataset', DATASETS, ids=["data", "prelim_data"])
+    def test_pull_nhsn_data_output(self, mock_create_backup, mock_socrata, dataset, caplog, params):
         now = time.time()
         # Mock Socrata client and its get method
         mock_client = MagicMock()
         mock_socrata.return_value = mock_client
-        mock_client.get.side_effect = [TEST_DATA,[]]
-
+        mock_client.get.side_effect = [dataset["test_data"],[]]
         mock_client.get_metadata.return_value = {"rowsUpdatedAt": now}
 
         backup_dir = params["common"]["backup_dir"]
         test_token = params["indicator"]["socrata_token"]
         custom_run = params["common"]["custom_run"]
-
         logger = get_structured_logger()
 
-        result = pull_nhsn_data(test_token, backup_dir, custom_run, issue_date=None, logger=logger)
+        expected_df = pd.read_csv(dataset["expected_data"])
+
+        result = pull_nhsn_data(test_token, backup_dir, custom_run, issue_date=None, logger=logger, preliminary=dataset["prelim_flag"])
         mock_create_backup.assert_called_once()
 
-        expected_columns = set(TYPE_DICT.keys())
+        expected_columns = set(expected_df.columns)
         assert set(result.columns) == expected_columns
 
         for column in list(result.columns):
             # some states don't report confirmed admissions rsv
-            if column == "confirmed_admissions_rsv_ew":
+            if column == "confirmed_admissions_rsv_ew" and not dataset["prelim_flag"]:
+                continue
+            if column == "confirmed_admissions_rsv_ew_prelim" and dataset["prelim_flag"]:
                 continue
             assert result[column].notnull().all(), f"{column} has rogue NaN"
 
+        expected_df = expected_df.astype(dataset["type_dict"])
 
-    @patch("delphi_nhsn.pull.Socrata")
-    def test_pull_nhsn_data_backup(self, mock_socrata, caplog, params):
-        now = time.time()
-        # Mock Socrata client and its get method
-        mock_client = MagicMock()
-        mock_socrata.return_value = mock_client
-        mock_client.get.side_effect = [TEST_DATA, []]
-
-        mock_client.get_metadata.return_value = {"rowsUpdatedAt": now}
+        pd.testing.assert_frame_equal(expected_df, result)
 
-        today = pd.Timestamp.today().strftime("%Y%m%d")
-        backup_dir = params["common"]["backup_dir"]
-        custom_run = params["common"]["custom_run"]
-        test_token = params["indicator"]["socrata_token"]
-
-        # Load test data
-        expected_data = pd.DataFrame(TEST_DATA)
-
-        logger = get_structured_logger()
-        # Call function with test token
-        pull_nhsn_data(test_token, backup_dir, custom_run, issue_date=None, logger=logger)
-
-        # Check logger used:
-        assert "Backup file created" in caplog.text
-
-        # Check that backup file was created
-        backup_files = glob.glob(f"{backup_dir}/{today}*")
-        assert len(backup_files) == 2, "Backup file was not created"
-
-        for backup_file in backup_files:
-            if backup_file.endswith(".csv.gz"):
-                dtypes = expected_data.dtypes.to_dict()
-                actual_data = pd.read_csv(backup_file, dtype=dtypes)
-            else:
-                actual_data = pd.read_parquet(backup_file)
-            pd.testing.assert_frame_equal(expected_data, actual_data)
-
-        # clean up
-        for file in backup_files:
-            os.remove(file)
 
     @patch("delphi_nhsn.pull.Socrata")
-    @patch("delphi_nhsn.pull.create_backup_csv")
-    def test_pull_prelim_nhsn_data_output(self, mock_create_backup, mock_socrata, caplog, params):
+    @pytest.mark.parametrize('dataset', DATASETS, ids=["data", "prelim_data"])
+    def test_pull_nhsn_data_backup(self, mock_socrata, dataset, caplog, params):
         now = time.time()
         # Mock Socrata client and its get method
         mock_client = MagicMock()
         mock_socrata.return_value = mock_client
-        mock_client.get.side_effect = [PRELIM_TEST_DATA, []]
+        mock_client.get.side_effect = [dataset["test_data"], []]
 
         mock_client.get_metadata.return_value = {"rowsUpdatedAt": now}
 
-        backup_dir = params["common"]["backup_dir"]
-        test_token = params["indicator"]["socrata_token"]
-        custom_run = params["common"]["custom_run"]
-
-        logger = get_structured_logger()
-
-        result = pull_preliminary_nhsn_data(test_token, backup_dir, custom_run, issue_date=None, logger=logger)
-        mock_create_backup.assert_called_once()
-
-        expected_columns = set(PRELIM_TYPE_DICT.keys())
-        assert set(result.columns) == expected_columns
-
-        for column in list(result.columns):
-            # some states don't report confirmed admissions rsv
-            if column == "confirmed_admissions_rsv_ew_prelim":
-                continue
-            assert result[column].notnull().all(), f"{column} has rogue NaN"
-    @patch("delphi_nhsn.pull.Socrata")
-    def test_pull_prelim_nhsn_data_backup(self, mock_socrata, caplog, params):
-        now = time.time()
-        # Mock Socrata client and its get method
-        mock_client = MagicMock()
-        mock_socrata.return_value = mock_client
-        mock_client.get.side_effect = [PRELIM_TEST_DATA, []]
-
-        mock_client.get_metadata.return_value = {"rowsUpdatedAt": now}
         today = pd.Timestamp.today().strftime("%Y%m%d")
         backup_dir = params["common"]["backup_dir"]
         custom_run = params["common"]["custom_run"]
         test_token = params["indicator"]["socrata_token"]
 
         # Load test data
-        expected_data = pd.DataFrame(PRELIM_TEST_DATA)
+        expected_data = pd.DataFrame(dataset["test_data"])
 
         logger = get_structured_logger()
         # Call function with test token
-        pull_preliminary_nhsn_data(test_token, backup_dir, custom_run, issue_date=None, logger=logger)
+        pull_nhsn_data(test_token, backup_dir, custom_run, issue_date=None, logger=logger, preliminary=dataset["prelim_flag"])
 
         # Check logger used:
         assert "Backup file created" in caplog.text
@@ -212,6 +156,7 @@ def test_pull_prelim_nhsn_data_backup(self, mock_socrata, caplog, params):
         for file in backup_files:
             os.remove(file)
 
+
     @pytest.mark.parametrize('dataset', DATASETS, ids=["data", "prelim_data"])
     @pytest.mark.parametrize("updatedAt", [time.time(), time.time() - 172800], ids=["updated", "stale"])
     @patch("delphi_nhsn.pull.Socrata")