Test for blank tract removal (#1848)

usds · Sep 15, 2022 · b8ac68d · b8ac68d
1 parent 665d136
commit b8ac68d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/...pipeline/data_pipeline/tests/sources/dot_travel_composite/data/Shapefile_and_Metadata.zip b/...pipeline/data_pipeline/tests/sources/dot_travel_composite/data/Shapefile_and_Metadata.zip
diff --git a/data/data-pipeline/data_pipeline/tests/sources/dot_travel_composite/test_etl.py b/data/data-pipeline/data_pipeline/tests/sources/dot_travel_composite/test_etl.py
@@ -21,14 +21,25 @@ def setup_method(self, _method, filename=__file__):
         """
         super().setup_method(_method=_method, filename=filename)
 
-    def test_extract_produces_valid_data(self, snapshot, mock_etl, mock_paths):
+    def test_extract_produces_valid_data(self, mock_etl, mock_paths):
         etl = self._setup_etl_instance_and_run_extract(
             mock_etl=mock_etl,
             mock_paths=mock_paths,
         )
-        tmp_df = gpd.read_file(
+        df = gpd.read_file(
             etl.get_tmp_path() / self._SAMPLE_DATA_FILE_NAME,
             dtype={etl.GEOID_TRACT_FIELD_NAME: str},
         )
-        assert tmp_df.shape[0] >= 15
-        assert tmp_df.shape[1] >= 86
+        assert df.shape[0] == 30
+        assert df.shape[1] == 86
+
+    def test_transform_removes_blank_tracts(self, mock_etl, mock_paths):
+        etl: TravelCompositeETL = self._setup_etl_instance_and_run_extract(
+            mock_etl=mock_etl,
+            mock_paths=mock_paths,
+        )
+        etl.transform()
+        etl.load()
+        df = etl.get_data_frame()
+        assert df.shape[0] == 15
+        assert df.shape[1] == len(etl.COLUMNS_TO_KEEP)