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)
   }
 }