Strengthen assertion to check partitioning

apache · May 10, 2015 · 8f5061a · 8f5061a
1 parent 01afc74
commit 8f5061a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/core/src/test/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleWriterSuite.java
@@ -188,7 +188,9 @@ private List<Tuple2<Object, Object>> readRecordsFromFile() throws IOException {
           new LimitedInputStream(in, partitionSize));
         Iterator<Tuple2<Object, Object>> records = recordsStream.asKeyValueIterator();
         while (records.hasNext()) {
-          recordsList.add(records.next());
+          Tuple2<Object, Object> record = records.next();
+          Assert.assertEquals(i, hashPartitioner.getPartition(record._1()));
+          recordsList.add(record);
         }
         recordsStream.close();
         startOffset += partitionSize;