Skip to content

Commit fcfac36

Browse files
author
Louiszr
committed
Made save load tests more robust to new params
1 parent ba994fd commit fcfac36

File tree

1 file changed

+8
-0
lines changed

1 file changed

+8
-0
lines changed

python/pyspark/ml/tests/test_tuning.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,10 @@ def test_save_load_trained_model(self):
230230
]:
231231
self.assertEqual(param(cvModel), param(loadedCvModel))
232232

233+
self.assertTrue(all(
234+
loadedCvModel.isSet(param) for param in loadedCvModel.params
235+
))
236+
233237
def test_save_load_simple_estimator(self):
234238
temp_path = tempfile.mkdtemp()
235239
dataset = self.spark.createDataFrame(
@@ -602,6 +606,10 @@ def test_save_load_trained_model(self):
602606
]:
603607
self.assertEqual(param(tvsModel), param(loadedTvsModel))
604608

609+
self.assertTrue(all(
610+
loadedTvsModel.isSet(param) for param in loadedTvsModel.params
611+
))
612+
605613
def test_save_load_simple_estimator(self):
606614
# This tests saving and loading the trained model only.
607615
# Save/load for TrainValidationSplit will be added later: SPARK-13786

0 commit comments

Comments
 (0)