diff --git a/core/src/test/scala/spark/RDDSuite.scala b/core/src/test/scala/spark/RDDSuite.scala index 03aa2845f46309b20e701a13c0edeef77d35e610..ef74c99246728a2b4857914c26af9ac43cdc997b 100644 --- a/core/src/test/scala/spark/RDDSuite.scala +++ b/core/src/test/scala/spark/RDDSuite.scala @@ -173,7 +173,10 @@ class RDDSuite extends FunSuite with BeforeAndAfter { sc = new SparkContext("local", "test") val data = sc.parallelize(1 to 10, 10) // Note that split number starts from 0, so > 8 means only 10th partition left. - val prunedData = data.pruneSplits(splitNum => splitNum > 8).collect - assert(prunedData.size == 1 && prunedData(0) == 10) + val prunedRdd = data.pruneSplits(splitNum => splitNum > 8) + assert(prunedRdd.splits.size == 1) + val prunedData = prunedRdd.collect + assert(prunedData.size == 1) + assert(prunedData(0) == 10) } }