Use withSystemProperty in ShuffleSuite

apache · Dec 19, 2014 · 51aa870 · 51aa870
1 parent 60a63a1
commit 51aa870
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/core/src/test/scala/org/apache/spark/ShuffleSuite.scala b/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
@@ -35,8 +35,7 @@ abstract class ShuffleSuite extends FunSuite with Matchers with LocalSparkContex
   conf.set("spark.test.noStageRetry", "true")
 
   test("groupByKey without compression") {
-    try {
-      System.setProperty("spark.shuffle.compress", "false")
+    TestUtils.withSystemProperty("spark.shuffle.compress", "false") {
       sc = new SparkContext("local", "test", conf)
       val pairs = sc.parallelize(Array((1, 1), (1, 2), (1, 3), (2, 1)), 4)
       val groups = pairs.groupByKey(4).collect()
@@ -45,8 +44,6 @@ abstract class ShuffleSuite extends FunSuite with Matchers with LocalSparkContex
       assert(valuesFor1.toList.sorted === List(1, 2, 3))
       val valuesFor2 = groups.find(_._1 == 2).get._2
       assert(valuesFor2.toList.sorted === List(1))
-    } finally {
-      System.setProperty("spark.shuffle.compress", "true")
     }
   }