From 11ffdf115e8151e9dfdc48468d7ff9f0c9d792c9 Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Wed, 1 Jul 2015 17:18:42 +0900
Subject: [PATCH] Use `===` and the variable

---
 .../test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala | 2 +-
 1 file changed, 1 insertion(+), 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
index 6052d77475694..424d32437a8dc 100644
--- 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))
   }
 }