Fix Java 8 cogroup tests for the new API

witgo · Apr 8, 2014 · 2d06e10 · 2d06e10
1 parent 11e730c
commit 2d06e10
Showing 1 changed file with 15 additions and 5 deletions.
diff --git a/extras/java8-tests/src/test/java/org/apache/spark/Java8APISuite.java b/extras/java8-tests/src/test/java/org/apache/spark/Java8APISuite.java
@@ -60,6 +60,16 @@ public void tearDown() {
     System.clearProperty("spark.driver.port");
   }
 
+  private int iterableSize(Iterable<?> a) {
+    int count = 0;
+    Iterator aItr = a.iterator();
+    while (aItr.hasNext()) {
+      aItr.next();
+      count++;
+    }
+    return count;
+  }
+
   @Test
   public void foreachWithAnonymousClass() {
     foreachCalls = 0;
@@ -85,15 +95,15 @@ public void foreach() {
   public void groupBy() {
     JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 1, 2, 3, 5, 8, 13));
     Function<Integer, Boolean> isOdd = x -> x % 2 == 0;
-    JavaPairRDD<Boolean, List<Integer>> oddsAndEvens = rdd.groupBy(isOdd);
+    JavaPairRDD<Boolean, Iterable<Integer>> oddsAndEvens = rdd.groupBy(isOdd);
     Assert.assertEquals(2, oddsAndEvens.count());
-    Assert.assertEquals(2, oddsAndEvens.lookup(true).get(0).size());  // Evens
-    Assert.assertEquals(5, oddsAndEvens.lookup(false).get(0).size()); // Odds
+    Assert.assertEquals(2, iterableSize(oddsAndEvens.lookup(true).get(0)));  // Evens
+    Assert.assertEquals(5, iterableSize(oddsAndEvens.lookup(false).get(0))); // Odds
 
     oddsAndEvens = rdd.groupBy(isOdd, 1);
     Assert.assertEquals(2, oddsAndEvens.count());
-    Assert.assertEquals(2, oddsAndEvens.lookup(true).get(0).size());  // Evens
-    Assert.assertEquals(5, oddsAndEvens.lookup(false).get(0).size()); // Odds
+    Assert.assertEquals(2, iterableSize(oddsAndEvens.lookup(true).get(0)));  // Evens
+    Assert.assertEquals(5, iterableSize(oddsAndEvens.lookup(false).get(0))); // Odds
   }
 
   @Test