Use === and the variable

rxin · Jul 1, 2015 · 11ffdf1 · 11ffdf1
1 parent 220a176
commit 11ffdf1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
@@ -106,7 +106,7 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(transformed.columns === Array("features", predictionColName))
     val clusters = transformed.select(predictionColName)
         .map(row => row.apply(0)).distinct().collect().toSet
-    assert(clusters.size == 5)
+    assert(clusters.size === k)
     assert(clusters === Set(0, 1, 2, 3, 4))
   }
 }