2129 fix hrr for nssp

nssp/delphi_nssp/constants.py

@@ -1,4 +1,5 @@
 """Registry for variations."""
+DATASET_ID = "rdmq-nq56"
 
 GEOS = [
     "hrr",

nssp/delphi_nssp/pull.py

@@ -13,7 +13,7 @@
 from delphi_utils import create_backup_csv
 from sodapy import Socrata
 
-from .constants import NEWLINE, SIGNALS, SIGNALS_MAP, TYPE_DICT
+from .constants import NEWLINE, SIGNALS, SIGNALS_MAP, TYPE_DICT, DATASET_ID
 
 
 def print_callback(remote_file_name, logger, bytes_so_far, bytes_total, progress_chunks):
@@ -148,7 +148,7 @@ def pull_nssp_data(
         Dataframe as described above.
     """
     if not custom_run:
-        socrata_results = pull_with_socrata_api(socrata_token, "rdmq-nq56")
+        socrata_results = pull_with_socrata_api(socrata_token, DATASET_ID)
         df_ervisits = pd.DataFrame.from_records(socrata_results)
         create_backup_csv(df_ervisits, backup_dir, custom_run, logger=logger)
         logger.info("Number of records grabbed", num_records=len(df_ervisits), source="Socrata API")

nssp/delphi_nssp/run.py

@@ -117,11 +117,14 @@ def run_module(params, logger=None):
                     lambda x: us.states.lookup(x).abbr.lower() if us.states.lookup(x) else "dc"
                 )
             elif geo == "hrr":
+                df = df[df["county"] != "All"]
                 df = df[["fips", "val", "timestamp"]]
-                # fips -> hrr has a weighted version
-                df = geo_mapper.replace_geocode(df, "fips", "hrr")
-                df = df.rename(columns={"hrr": "geo_id"})
+                df = geo_mapper.add_population_column(df, geocode_type="fips", geocode_col="fips")
+                df = geo_mapper.add_geocode(df, "fips", "hrr", from_col="fips", new_col="geo_id")
+                df = geo_mapper.aggregate_by_weighted_sum(df, "geo_id", "val", "timestamp", "population")
+                df = df.rename(columns={"weighted_val": "val"})
             elif geo == "msa":
+                df = df[df["county"] != "All"]
                 df = df[["fips", "val", "timestamp"]]
                 # fips -> msa doesn't have a weighted version, so we need to add columns and sum ourselves
                 df = geo_mapper.add_population_column(df, geocode_type="fips", geocode_col="fips")
@@ -139,6 +142,7 @@ def run_module(params, logger=None):
             else:
                 df = df[df["county"] != "All"]
                 df["geo_id"] = df["fips"]
+
             # add se, sample_size, and na codes
             missing_cols = set(CSV_COLS) - set(df.columns)
             df = add_needed_columns(df, col_names=list(missing_cols))

nssp/tests/conftest.py

@@ -0,0 +1,89 @@
+import copy
+import json
+import time
+from unittest.mock import patch, MagicMock
+
+import pytest
+from pathlib import Path
+
+from delphi_nssp.run import run_module
+from delphi_nssp.constants import DATASET_ID
+
+TEST_DIR = Path(__file__).parent
+
+# test data generated with following url with socrata:
+# https://data.cdc.gov/resource/rdmq-nq56.json?$where=week_end >= '2022-10-01T00:00:00.000' AND week_end <= '2022-10-20T00:00:00.000'
+
+with open(f"{TEST_DIR}/test_data/page.json", "r") as f:
+    TEST_DATA = json.load(f)
+
+with open(f"{TEST_DIR}/test_data/page_100_hrr.json", "r") as f:
+    HRR_TEST_DATA = json.load(f)
+
+@pytest.fixture(scope="session")
+def params():
+    params = {
+        "common": {
+            "export_dir": f"{TEST_DIR}/receiving",
+            "log_filename": f"{TEST_DIR}/test.log",
+            "backup_dir": f"{TEST_DIR}/test_raw_data_backups",
+            "custom_run": False
+        },
+        "indicator": {
+            "wip_signal": True,
+            "export_start_date": "2020-08-01",
+            "static_file_dir": "./static",
+            "socrata_token": "test_token"
+        },
+        "validation": {
+            "common": {
+                "span_length": 14,
+                "min_expected_lag": {"all": "3"},
+                "max_expected_lag": {"all": "4"},
+            }
+        }
+    }
+    return copy.deepcopy(params)
+
+@pytest.fixture
+def params_w_patch(params):
+    params_copy = copy.deepcopy(params)
+    params_copy["common"]["custom_run"] = True
+    params_copy["patch"] = {
+            "patch_dir": f"{TEST_DIR}/patch_dir",
+            "source_dir": "test_source_dir",
+            "source_host": "host",
+            "user": "test_user",
+            "start_issue": "2023-01-01",
+            "end_issue": "2023-01-03",
+        }
+
+    return params_copy
+
+@pytest.fixture(scope="function")
+def run_as_module(params):
+    with patch('sodapy.Socrata.get') as mock_get:
+        def side_effect(*args, **kwargs):
+            if kwargs['offset'] == 0:
+                if DATASET_ID in args[0]:
+                    return TEST_DATA
+            else:
+                return []
+        mock_get.side_effect = side_effect
+        run_module(params)
+
+
+
+@pytest.fixture(scope="function")
+def run_as_module_hrr(params):
+    with patch('sodapy.Socrata.get') as mock_get, \
+         patch('delphi_nssp.run.GEOS', ["hrr"]):
+        def side_effect(*args, **kwargs):
+            if kwargs['offset'] == 0:
+                if DATASET_ID in args[0]:
+                    return HRR_TEST_DATA
+            else:
+                return []
+        mock_get.side_effect = side_effect
+        run_module(params)
+