fixed

apache · May 29, 2018 · e197460 · e197460
1 parent b9d8dd3
commit e197460
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/python/pyspark/tests.py b/python/pyspark/tests.py
@@ -545,13 +545,15 @@ def test_tc_on_driver(self):
 
     def test_get_local_property(self):
         """Verify that local properties set on the driver are available in TaskContext."""
-        self.sc.setLocalProperty("testkey", "testvalue")
-        rdd = self.sc.parallelize(range(1), 1)
-        prop1 = rdd.map(lambda x: TaskContext.get().getLocalProperty("testkey")).collect()[0]
-        self.assertEqual(prop1, "testkey")
-        prop2 = rdd.map(lambda x: TaskContext.get().getLocalProperty("otherkey")).collect()[0]
-        self.assertTrue(prop2 is None)
-
+        try:
+            self.sc.setLocalProperty("testkey", "testvalue")
+            rdd = self.sc.parallelize(range(1), 1)
+            prop1 = rdd.map(lambda x: TaskContext.get().getLocalProperty("testkey")).collect()[0]
+            self.assertEqual(prop1, "testvalue")
+            prop2 = rdd.map(lambda x: TaskContext.get().getLocalProperty("otherkey")).collect()[0]
+            self.assertTrue(prop2 is None)
+        finally:
+            self.sc.setLocalProperty("testkey", None)
 
 class RDDTests(ReusedPySparkTestCase):