Python unit tests check read pipeline stages

Signed-off-by: Jason T. Brown <[email protected]>

Author: vpipkt

pyrasterframes/src/main/python/tests/ExploderTests.py

@@ -68,3 +68,4 @@ def test_tile_exploder_read_write(self):
 
         read_pipe = PipelineModel.load(path)
         self.assertEqual(len(read_pipe.stages), 2)
+        self.assertTrue(isinstance(read_pipe.stages[0], TileExploder))

pyrasterframes/src/main/python/tests/NoDataFilterTests.py

@@ -37,12 +37,15 @@ def test_no_data_filter_read_write(self):
         df = self.spark.read.raster(self.img_uri) \
             .select(rf_tile_mean('proj_raster').alias('mean'))
 
-        ndf = NoDataFilter().setInputCols(['mean'])
-        assembler = VectorAssembler().setInputCols(['mean'])
+        input_cols = ['mean']
+        ndf = NoDataFilter().setInputCols(input_cols)
+        assembler = VectorAssembler().setInputCols(input_cols)
 
         pipe = Pipeline().setStages([ndf, assembler])
 
         pipe.fit(df).write().overwrite().save(path)
 
         read_pipe = PipelineModel.load(path)
         self.assertEqual(len(read_pipe.stages), 2)
+        actual_stages_ndf = read_pipe.stages[0].getInputCols()
+        self.assertEqual(actual_stages_ndf, input_cols)