From c0c20647c50ed59f1cbdc971e12d725488a87b8e Mon Sep 17 00:00:00 2001 From: matt bowen Date: Fri, 2 Sep 2022 11:24:57 -0400 Subject: [PATCH] Apply code-review feedback (#1835) --- .../data_pipeline/tests/score/test_output.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/data/data-pipeline/data_pipeline/tests/score/test_output.py b/data/data-pipeline/data_pipeline/tests/score/test_output.py index 0ad06b616..3d8b743ae 100644 --- a/data/data-pipeline/data_pipeline/tests/score/test_output.py +++ b/data/data-pipeline/data_pipeline/tests/score/test_output.py @@ -263,16 +263,11 @@ def test_data_sources( suffixes=(final, f"_{data_source_name}"), how="left", ) - data_source_columns = [ - f"{col}_{data_source_name}" - for col in data_source.columns - if (col != GEOID_TRACT_FIELD_NAME and col in final_score_df.columns) - ] - final_columns = [ - f"{col}{final}" - for col in data_source.columns - if (col != GEOID_TRACT_FIELD_NAME and col in final_score_df.columns) - ] + core_cols = data_source.columns.intersection( + final_score_df.columns + ).drop(GEOID_TRACT_FIELD_NAME) + data_source_columns = [f"{col}_{data_source_name}" for col in core_cols] + final_columns = [f"{col}{final}" for col in core_cols] assert np.all(df[df.MERGE == "left_only"][final_columns].isna()) df = df[df.MERGE == "both"] assert (