1 files changed, 7 insertions, 4 deletions
diff --git a/mllib/src/test/scala/spark/mllib/clustering/KMeansSuite.scala b/mllib/src/test/scala/spark/mllib/clustering/KMeansSuite.scala
index ae7cf57c42..cb096f39a9 100644
--- a/mllib/src/test/scala/spark/mllib/clustering/KMeansSuite.scala
+++ b/mllib/src/test/scala/spark/mllib/clustering/KMeansSuite.scala
@@ -21,6 +21,8 @@ class KMeansSuite extends FunSuite with BeforeAndAfterAll {
 
   val EPSILON = 1e-4
 
+  import KMeans.{RANDOM, K_MEANS_PARALLEL}
+
   def prettyPrint(point: Array[Double]): String = point.mkString("(", ", ", ")")
 
   def prettyPrint(points: Array[Array[Double]]): String = {
@@ -82,10 +84,11 @@ class KMeansSuite extends FunSuite with BeforeAndAfterAll {
     model = KMeans.train(data, k=1, maxIterations=1, runs=5)
     assertSetsEqual(model.clusterCenters, Array(Array(1.0, 3.0, 4.0)))
 
-    model = KMeans.train(data, k=1, maxIterations=1, runs=1, initializationMode="random")
+    model = KMeans.train(data, k=1, maxIterations=1, runs=1, initializationMode=RANDOM)
     assertSetsEqual(model.clusterCenters, Array(Array(1.0, 3.0, 4.0)))
 
-    model = KMeans.train(data, k=1, maxIterations=1, runs=1, initializationMode="k-means||")
+    model = KMeans.train(
+      data, k=1, maxIterations=1, runs=1, initializationMode=K_MEANS_PARALLEL)
     assertSetsEqual(model.clusterCenters, Array(Array(1.0, 3.0, 4.0)))
   }
 
@@ -115,10 +118,10 @@ class KMeansSuite extends FunSuite with BeforeAndAfterAll {
     model = KMeans.train(data, k=1, maxIterations=1, runs=5)
     assertSetsEqual(model.clusterCenters, Array(Array(1.0, 3.0, 4.0)))
 
-    model = KMeans.train(data, k=1, maxIterations=1, runs=1, initializationMode="random")
+    model = KMeans.train(data, k=1, maxIterations=1, runs=1, initializationMode=RANDOM)
     assertSetsEqual(model.clusterCenters, Array(Array(1.0, 3.0, 4.0)))
 
-    model = KMeans.train(data, k=1, maxIterations=1, runs=1, initializationMode="k-means||")
+    model = KMeans.train(data, k=1, maxIterations=1, runs=1, initializationMode=K_MEANS_PARALLEL)
     assertSetsEqual(model.clusterCenters, Array(Array(1.0, 3.0, 4.0)))
   }