Apply code-review feedback (#1835)

usds · Sep 2, 2022 · c0c2064 · c0c2064
1 parent d90292b
commit c0c2064
Showing 1 changed file with 5 additions and 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
@@ -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 (