Better test

apache · Jan 11, 2017 · 3dea44f · 3dea44f
1 parent 30ed14f
commit 3dea44f
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaUDFSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaUDFSuite.java
@@ -18,6 +18,7 @@
 package test.org.apache.spark.sql;
 
 import java.io.Serializable;
+import java.util.List;
 
 import org.junit.After;
 import org.junit.Assert;
@@ -120,7 +121,12 @@ public Long call(Long i) {
     }, DataTypes.LongType);
 
     spark.range(10).toDF("x").createOrReplaceTempView("tmp");
-    Row result = spark.sql("SELECT inc(x) FROM tmp GROUP BY inc(x)").head();
-    Assert.assertEquals(7, result.getLong(0));
+    List<Row> results = spark.sql("SELECT inc(x) FROM tmp GROUP BY inc(x)").collectAsList();
+    Assert.assertEquals(10, results.size());
+    long sum = 0;
+    for (Row result : results) {
+      sum += result.getLong(0);
+    }
+    Assert.assertEquals(55, sum);
   }
 }