222 works; still some warnings but tests and lint passing

SatcherInstitute · May 10, 2024 · 9ea70be · 9ea70be
1 parent 6e9406f
commit 9ea70be
Show file tree

Hide file tree

Showing 7 changed files with 195 additions and 255 deletions.
diff --git a/airflow/dags/graphql_ahr.py b/airflow/dags/graphql_ahr.py
@@ -1,5 +1,5 @@
-from airflow import DAG  # type: ignore
-from airflow.utils.dates import days_ago  # type: ignore
+from airflow import DAG  # pylint: disable = no-name-in-module
+from airflow.utils.dates import days_ago  # pylint: disable = no-name-in-module
 import util
 
 _GRAPHQL_AHR_WORKFLOW_ID = 'GRAPHQL_AHR_DATA'

diff --git a/python/ingestion/graphql_ahr_utils.py b/python/ingestion/graphql_ahr_utils.py
@@ -110,7 +110,7 @@ def fetch_ahr_data_from_graphql():
         """
         variables = json.loads(variables_str)
         graphql_request = {'query': query, 'variables': variables}
-        response = requests.post(graphql_url, json=graphql_request, headers=headers)
+        response = requests.post(graphql_url, json=graphql_request, headers=headers, timeout=10)
 
         if response.status_code == 200:
             # Collect each successful responses

diff --git a/python/tests/data/graphql_ahr/golden_data/age_national_current.csv b/python/tests/data/graphql_ahr/golden_data/age_national_current.csv
@@ -7,7 +7,7 @@ United States,00,35-44,,,,,,,,,,,,,18.1
 United States,00,45-54,,,,,,,,,,,,,18.2
 United States,00,55-64,,,,,,,,,,,,,17.0
 United States,00,65+,,8600.0,3.4,20600.0,7100.0,12300.0,14600.0,22600.0,7000.0,8500.0,,,16.9
-United States,00,65-74,,,,,,,,,,,1482.0,,15.3
+United States,00,65-74,,,,,,,,,,,,,15.3
 United States,00,75-84,,,,,,,,,,,,,19.6
 United States,00,85+,,,,,,,,,,,,,22.4
 United States,00,All,,9800.0,8.8,8000.0,3000.0,6200.0,20500.0,10900.0,17300.0,14700.0,2681.0,12.0,14.5
diff --git a/python/tests/data/graphql_ahr/golden_data/race_and_ethnicity_state_current.csv b/python/tests/data/graphql_ahr/golden_data/race_and_ethnicity_state_current.csv
diff --git a/python/tests/data/graphql_ahr/golden_data/sex_national_current.csv b/python/tests/data/graphql_ahr/golden_data/sex_national_current.csv
@@ -1,4 +1,4 @@
 state_name,state_fips,sex,voter_participation_pct_rate,asthma_per_100k,avoided_care_pct_rate,cardiovascular_diseases_per_100k,chronic_kidney_disease_per_100k,copd_per_100k,depression_per_100k,diabetes_per_100k,excessive_drinking_per_100k,frequent_mental_distress_per_100k,preventable_hospitalizations_per_100k,non_medical_drug_use_per_100k,suicide_per_100k
-United States,00,Male,,6800.0,9.1,9400.0,2900.0,5700.0,14300.0,11800.0,20800.0,12000.0,2598.0,,23.7
-United States,00,Female,,12400.0,10.6,7300.0,3300.0,7100.0,24400.0,11000.0,13400.0,17200.0,2756.0,,5.9
+United States,00,Male,,6800.0,9.1,9400.0,2900.0,5700.0,14300.0,11800.0,20800.0,12000.0,,,23.7
+United States,00,Female,,12400.0,10.6,7300.0,3300.0,7100.0,24400.0,11000.0,13400.0,17200.0,,,5.9
 United States,00,All,,9800.0,8.8,8000.0,3000.0,6200.0,20500.0,10900.0,17300.0,14700.0,2681.0,12.0,14.5
diff --git a/python/tests/datasources/test_decia_2020_territory_population.py b/python/tests/datasources/test_decia_2020_territory_population.py
@@ -12,6 +12,7 @@
 
 def _load_csv_as_df_from_data_dir(*args, **kwargs):
     directory, filename = args
+    print("kwargs:", kwargs)
     df = pd.read_csv(os.path.join(TEST_DIR, directory, filename))
     return df
 
@@ -64,6 +65,9 @@ def testGenerateRaceTerritory(
     assert table_name == "by_race_and_ethnicity_territory_state_level"
     expected_df = pd.read_csv(os.path.join(GOLDEN_DIR, f'{table_name}.csv'), index_col=False, dtype=dtypes)
 
+    df = df.sort_values(by=['state_fips', 'race_category_id']).reset_index(drop=True)
+    expected_df = expected_df.sort_values(by=['state_fips', 'race_category_id']).reset_index(drop=True)
+
     assert_frame_equal(df, expected_df, check_dtype=False, check_like=True)
 
 

diff --git a/python/tests/datasources/test_graphql_ahr.py b/python/tests/datasources/test_graphql_ahr.py
@@ -22,14 +22,15 @@
 
 
 def _fetch_ahr_data_from_graphql():
+    print("MOCK - AHR GraphQL API response")
     with open(EXPECTED_AHR_API_RESPONSE_DATA, 'r', encoding='utf-8') as file:
         data = json.load(file)
 
     return data
 
 
 @mock.patch('ingestion.gcs_to_bq_util.add_df_to_bq', return_value=None)
-@mock.patch('ingestion.graphql_ahr_utils.fetch_ahr_data_from_graphql', side_effect=_fetch_ahr_data_from_graphql)
+@mock.patch('datasources.graphql_ahr.fetch_ahr_data_from_graphql', side_effect=_fetch_ahr_data_from_graphql)
 def testWriteToBqAgeNational(_mock_fetch: mock.MagicMock, mock_add_df_to_bq: mock.MagicMock):
     datasource = GraphQlAHRData()
     datasource.write_to_bq('dataset', 'gcs_bucket', demographic=AGE, geographic=NATIONAL_LEVEL)
@@ -44,7 +45,7 @@ def testWriteToBqAgeNational(_mock_fetch: mock.MagicMock, mock_add_df_to_bq: moc
 
 
 @mock.patch('ingestion.gcs_to_bq_util.add_df_to_bq', return_value=None)
-@mock.patch('ingestion.graphql_ahr_utils.fetch_ahr_data_from_graphql', side_effect=_fetch_ahr_data_from_graphql)
+@mock.patch('datasources.graphql_ahr.fetch_ahr_data_from_graphql', side_effect=_fetch_ahr_data_from_graphql)
 def testWriteToBqRaceNational(_mock_fetch: mock.MagicMock, mock_add_df_to_bq: mock.MagicMock):
     datasource = GraphQlAHRData()
     datasource.write_to_bq('dataset', 'gcs_bucket', demographic=RACE_OR_HISPANIC_COL, geographic=STATE_LEVEL)
@@ -59,7 +60,7 @@ def testWriteToBqRaceNational(_mock_fetch: mock.MagicMock, mock_add_df_to_bq: mo
 
 
 @mock.patch('ingestion.gcs_to_bq_util.add_df_to_bq', return_value=None)
-@mock.patch('ingestion.graphql_ahr_utils.fetch_ahr_data_from_graphql', side_effect=_fetch_ahr_data_from_graphql)
+@mock.patch('datasources.graphql_ahr.fetch_ahr_data_from_graphql', side_effect=_fetch_ahr_data_from_graphql)
 def testWriteToBqSexNational(_mock_fetch: mock.MagicMock, mock_add_df_to_bq: mock.MagicMock):
     datasource = GraphQlAHRData()
     datasource.write_to_bq('dataset', 'gcs_bucket', demographic=SEX, geographic=NATIONAL_LEVEL)