TST: python test for setWeightCol

Author: facaiy

python/pyspark/ml/tests.py

@@ -1255,15 +1255,16 @@ def test_output_columns(self):
         output = model.transform(df)
         self.assertEqual(output.columns, ["label", "features", "prediction"])
 
-    def test_cache_weightCol_if_necessary(self):
+    def test_support_for_weightCol(self):
         df = self.spark.createDataFrame([(0.0, Vectors.dense(1.0, 0.8), 1.0),
                                          (1.0, Vectors.sparse(2, [], []), 1.0),
                                          (2.0, Vectors.dense(0.5, 0.5), 1.0)],
                                         ["label", "features", "weight"])
         lr = LogisticRegression(maxIter=5, regParam=0.01)
         ovr = OneVsRest(classifier=lr, weightCol="weight")
-        model = ovr.fit(df)
-        self.assertIsNotNone(model)
+        self.assertIsNotNone(ovr.fit(df))
+        ovr2 = OneVsRest(classifier=lr).setWeightCol("weight")
+        self.assertIsNotNone(ovr2.fit(df))
 
 
 class HashingTFTest(SparkSessionTestCase):