Add a test case for random initialization.

Also workaround a bug where double[][] class cast fails

1 files changed, 11 insertions, 0 deletions

diff --git a/mllib/src/test/scala/spark/mllib/clustering/JavaKMeansSuite.java b/mllib/src/test/scala/spark/mllib/clustering/JavaKMeansSuite.java
index f438a92fad..a3db0c0f6d 100644
--- a/mllib/src/test/scala/spark/mllib/clustering/JavaKMeansSuite.java
+++ b/mllib/src/test/scala/spark/mllib/clustering/JavaKMeansSuite.java
@@ -86,6 +86,10 @@ public class JavaKMeansSuite implements Serializable {
 
     JavaRDD<double[]> data = sc.parallelize(points, 2);
     KMeansModel model = KMeans.train(data.rdd(), 1, 1);
+    assertSetsEqual(model.clusterCenters(), expectedCenter);
+
+    model = KMeans.train(data.rdd(), 1, 1, 1, KMeans.RANDOM());
+    assertSetsEqual(model.clusterCenters(), expectedCenter);
   }
 
   @Test
@@ -100,5 +104,12 @@ public class JavaKMeansSuite implements Serializable {
     JavaRDD<double[]> data = sc.parallelize(points, 2);
     KMeansModel model = new KMeans().setK(1).setMaxIterations(5).run(data.rdd());
     assertSetsEqual(model.clusterCenters(), expectedCenter);
+
+    model = new KMeans().setK(1)
+                        .setMaxIterations(1)
+                        .setRuns(1)
+                        .setInitializationMode(KMeans.RANDOM())
+                        .run(data.rdd());
+    assertSetsEqual(model.clusterCenters(), expectedCenter);
   }
 }