Switch the comparisons.

rxin · Jul 17, 2015 · f397be4 · f397be4
1 parent ca78b7d
commit f397be4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/mllib/src/test/java/org/apache/spark/ml/clustering/JavaKMeansSuite.java b/mllib/src/test/java/org/apache/spark/ml/clustering/JavaKMeansSuite.java
@@ -63,9 +63,10 @@ public void fitAndTransform() {
     assertEquals(k, centers.length);
 
     DataFrame transformed = model.transform(dataset);
+    List<String> columns = Arrays.asList(transformed.columns());
     List<String> expectedColumns = Arrays.asList("features", "prediction");
-    for (String clm: transformed.columns()) {
-      assertTrue(expectedColumns.contains(clm));
+    for (String column: expectedColumns) {
+      columns.contains(column);
     }
   }
 }
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
@@ -103,8 +103,9 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(model.clusterCenters.length === k)
 
     val transformed = model.transform(dataset)
-    transformed.columns.foreach { column =>
-      Array("features", predictionColName).contains(column)
+    val expectedColumns = Array("features", predictionColName)
+    expectedColumns.foreach { column =>
+      transformed.columns.contains(column)
     }
     val clusters = transformed.select(predictionColName).map(_.get(0)).distinct().collect().toSet
     assert(clusters.size === k)