Make combo cases/deaths robust to source outages

- okay to return just one of the data frames if both aren't available
- do something sensible if neither one is available
- unit tests to verify all possible cases are handled appropriately

Fixes #330

Author: krivard

combo_cases_and_deaths/delphi_combo_cases_and_deaths/constants.py

@@ -18,4 +18,4 @@
     "state",
     "msa",
     "hrr",
-]
+]

combo_cases_and_deaths/delphi_combo_cases_and_deaths/run.py

@@ -9,46 +9,68 @@
 from datetime import date, timedelta, datetime
 from itertools import product
 import re
-import sys
 
 import covidcast
 import pandas as pd
 
-from delphi_utils import read_params, create_export_csv
-from .constants import *
-from .handle_wip_signal import *
+from delphi_utils import read_params
+from .constants import METRICS, SMOOTH_TYPES, SENSORS, GEO_RESOLUTIONS
+from .handle_wip_signal import add_prefix
 
 
 def check_not_none(data_frame, label, date_range):
     """Exit gracefully if a data frame we attempted to retrieve is empty"""
     if data_frame is None:
         print(f"{label} not available in range {date_range}")
-        sys.exit(1)
+        return False
+    return True
 
-def combine_usafacts_and_jhu(signal, geo, date_range):
+def maybe_append(df1, df2):
+    """
+    If both data frames are available, append them and return. Otherwise, return
+    whichever frame is not None.
+    """
+    if df1 is None:
+        return df2
+    if df2 is None:
+        return df1
+    return df1.append(df2)
+
+COLUMN_MAPPING = {"time_value": "timestamp",
+                  "geo_value": "geo_id",
+                  "value": "val",
+                  "stderr": "se",
+                  "sample_size": "sample_size"}
+def combine_usafacts_and_jhu(signal, geo, date_range, fetcher=covidcast.signal):
     """
     Add rows for PR from JHU signals to USA-FACTS signals
     """
-    usafacts_df = covidcast.signal("usa-facts", signal, date_range[0], date_range[1], geo)
-    jhu_df = covidcast.signal("jhu-csse", signal, date_range[0], date_range[1], geo)
-    check_not_none(usafacts_df, "USA-FACTS", date_range)
-    check_not_none(jhu_df, "JHU", date_range)
+    print("Fetching usa-facts...")
+    usafacts_df = fetcher("usa-facts", signal, date_range[0], date_range[1], geo)
+    print("Fetching jhu-csse...")
+    jhu_df = fetcher("jhu-csse", signal, date_range[0], date_range[1], geo)
+
+    if (not check_not_none(usafacts_df, "USA-FACTS", date_range)) and \
+       (geo not in ('state', 'county') or \
+        (not check_not_none(jhu_df, "JHU", date_range))):
+        return pd.DataFrame({}, columns=COLUMN_MAPPING.values())
 
     # State level
     if geo == 'state':
-        combined_df = usafacts_df.append(jhu_df[jhu_df["geo_value"] == 'pr'])
+        combined_df = maybe_append(
+            usafacts_df,
+            jhu_df if jhu_df is None else jhu_df[jhu_df["geo_value"] == 'pr'])
     # County level
     elif geo == 'county':
-        combined_df = usafacts_df.append(jhu_df[jhu_df["geo_value"] == '72000'])
+        combined_df = maybe_append(
+            usafacts_df,
+            jhu_df if jhu_df is None else jhu_df[jhu_df["geo_value"] == '72000'])
     # For MSA and HRR level, they are the same
     else:
         combined_df = usafacts_df
 
     combined_df = combined_df.drop(["direction"], axis=1)
-    combined_df = combined_df.rename({"time_value": "timestamp",
-                                      "geo_value": "geo_id",
-                                      "value": "val",
-                                      "stderr": "se"},
+    combined_df = combined_df.rename(COLUMN_MAPPING,
                                      axis=1)
     return combined_df
 
@@ -83,15 +105,12 @@ def sensor_signal(metric, sensor, smoother):
         sensor_name = "_".join([smoother, sensor])
     else:
         sensor_name = sensor
-    signal = "_".join([metric, sensor_name])
-    return sensor_name, signal
-
-def run_module():
-    """Produce a combined cases and deaths signal using data from JHU and USA Facts"""
-    variants = [tuple((metric, geo_res)+sensor_signal(metric, sensor, smoother))
-                for (metric, geo_res, sensor, smoother) in
-                product(METRICS, GEO_RESOLUTIONS, SENSORS, SMOOTH_TYPES)]
+    return sensor_name, "_".join([metric, sensor_name])
 
+def configure(variants):
+    """
+    Validate params file and set date range.
+    """
     params = read_params()
     params['export_start_date'] = date(*params['export_start_date'])
     yesterday = date.today() - timedelta(days=1)
@@ -112,30 +131,38 @@ def run_module():
         # create combined files for all of the historical reports
         params['date_range'] = [params['export_start_date'], yesterday]
     else:
-        pattern = re.compile(r'^\d{8}-\d{8}$')
-        match_res = re.findall(pattern, params['date_range'])
+        match_res = re.findall(re.compile(r'^\d{8}-\d{8}$'), params['date_range'])
         if len(match_res) == 0:
             raise ValueError(
                 "Invalid date_range parameter. Please choose from (new, all, yyyymmdd-yyyymmdd).")
         try:
             date1 = datetime.strptime(params['date_range'][:8], '%Y%m%d').date()
-        except ValueError:
-            raise ValueError("Invalid date_range parameter. Please check the first date.")
+        except ValueError as error:
+            raise ValueError(
+                "Invalid date_range parameter. Please check the first date.") from error
         try:
             date2 = datetime.strptime(params['date_range'][-8:], '%Y%m%d').date()
-        except ValueError:
-            raise ValueError("Invalid date_range parameter. Please check the second date.")
+        except ValueError as error:
+            raise ValueError(
+                "Invalid date_range parameter. Please check the second date.") from error
 
         #The the valid start date
         if date1 < params['export_start_date']:
             date1 = params['export_start_date']
         params['date_range'] = [date1, date2]
+    return params
 
-    for metric, geo_res, sensor_name, signal in variants:
 
-        df = combine_usafacts_and_jhu(signal, geo_res, extend_raw_date_range(params, sensor_name))
+def run_module():
+    """Produce a combined cases and deaths signal using data from JHU and USA Facts"""
+    variants = [tuple((metric, geo_res)+sensor_signal(metric, sensor, smoother))
+                for (metric, geo_res, sensor, smoother) in
+                product(METRICS, GEO_RESOLUTIONS, SENSORS, SMOOTH_TYPES)]
+    params = configure(variants)
+    for metric, geo_res, sensor_name, signal in variants:
+        df = combine_usafacts_and_jhu(signal, geo_res, extend_raw_date_range(params, sensor_name)) # pylint: disable=invalid-name
 
-        df = df.copy()
+        df = df.copy() # pylint: disable=invalid-name
         df["timestamp"] = pd.to_datetime(df["timestamp"])
         start_date = pd.to_datetime(params['export_start_date'])
         export_dir = params["export_dir"]
@@ -145,8 +172,7 @@ def run_module():
 
         signal_name = add_prefix([signal], wip_signal=params["wip_signal"], prefix="wip_")
         for date_ in dates:
-            export_fn = f'{date_.strftime("%Y%m%d")}_{geo_res}_' f"{signal_name[0]}.csv"
+            export_fn = f'{date_.strftime("%Y%m%d")}_{geo_res}_{signal_name[0]}.csv'
             df[df["timestamp"] == date_][["geo_id", "val", "se", "sample_size", ]].to_csv(
                 f"{export_dir}/{export_fn}", index=False, na_rep="NA"
             )
-

combo_cases_and_deaths/tests/test_run.py

@@ -2,11 +2,12 @@
 from itertools import product
 import pytest
 import unittest
+import pandas as pd
 
-from delphi_combo_cases_and_deaths.run import extend_raw_date_range, sensor_signal
+from delphi_combo_cases_and_deaths.run import extend_raw_date_range, sensor_signal, combine_usafacts_and_jhu, COLUMN_MAPPING
 from delphi_combo_cases_and_deaths.handle_wip_signal import add_prefix
 from delphi_utils import read_params
-from delphi_combo_cases_and_deaths.constants import *
+from delphi_combo_cases_and_deaths.constants import METRICS, SMOOTH_TYPES, SENSORS, GEO_RESOLUTIONS
 
 
 def test_issue_dates():
@@ -49,6 +50,33 @@ def test_handle_wip_signal():
     assert signal_names[0].startswith("wip_")
     assert all(not s.startswith("wip_") for s in signal_names[1:])
 
+def test_unstable_sources():
+    placeholder = lambda geo: pd.DataFrame(
+        [(date.today(),"pr" if geo == "state" else "72000",1,1,1,0)],
+        columns="time_value geo_value value stderr sample_size direction".split())
+    fetcher10 = lambda *x: placeholder(x[-1]) if x[0] == "usa-facts" else None
+    fetcher01 = lambda *x: placeholder(x[-1]) if x[0] == "jhu-csse" else None
+    fetcher11 = lambda *x: placeholder(x[-1])
+    fetcher00 = lambda *x: None
+
+    date_range = [date.today(), date.today()]
+
+    for geo in "state county msa".split():
+        for (fetcher, expected_size) in [
+                (fetcher00, 0),
+                (fetcher01, 0 if geo == "msa" else 1),
+                (fetcher10, 1),
+                (fetcher11, 1 if geo == "msa" else 2)
+        ]:
+            df = combine_usafacts_and_jhu("", geo, date_range, fetcher)
+            assert df.size == expected_size * len(COLUMN_MAPPING), f"""
+input for {geo}:
+{fetcher('usa-facts',geo)}
+{fetcher('jhu-csse',geo)}
+
+output:
+{df}
+"""
 
 class MyTestCase(unittest.TestCase):
     pass