From 866e6949df8cdba0a4da20a1045333e419ed0e12 Mon Sep 17 00:00:00 2001
From: Denny <dennybritz@gmail.com>
Date: Wed, 18 Jul 2012 13:09:50 -0700
Subject: Always destroy SparkContext in after block for the unit tests.

Conflicts:

	core/src/test/scala/spark/ShuffleSuite.scala
---
 bagel/src/test/scala/bagel/BagelSuite.scala | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

(limited to 'bagel/src')

diff --git a/bagel/src/test/scala/bagel/BagelSuite.scala b/bagel/src/test/scala/bagel/BagelSuite.scala
index 0eda80af64..5ac7f5d381 100644
--- a/bagel/src/test/scala/bagel/BagelSuite.scala
+++ b/bagel/src/test/scala/bagel/BagelSuite.scala
@@ -1,6 +1,6 @@
 package spark.bagel
 
-import org.scalatest.{FunSuite, Assertions}
+import org.scalatest.{FunSuite, Assertions, BeforeAndAfter}
 import org.scalatest.prop.Checkers
 import org.scalacheck.Arbitrary._
 import org.scalacheck.Gen
@@ -13,9 +13,16 @@ import spark._
 class TestVertex(val active: Boolean, val age: Int) extends Vertex with Serializable
 class TestMessage(val targetId: String) extends Message[String] with Serializable
 
-class BagelSuite extends FunSuite with Assertions {
+class BagelSuite extends FunSuite with Assertions with BeforeAndAfter{
+  
+  var sc: SparkContext = _
+  
+  after{
+    sc.stop()
+  }
+  
   test("halting by voting") {
-    val sc = new SparkContext("local", "test")
+    sc = new SparkContext("local", "test")
     val verts = sc.parallelize(Array("a", "b", "c", "d").map(id => (id, new TestVertex(true, 0))))
     val msgs = sc.parallelize(Array[(String, TestMessage)]())
     val numSupersteps = 5
@@ -26,11 +33,10 @@ class BagelSuite extends FunSuite with Assertions {
       }
     for ((id, vert) <- result.collect)
       assert(vert.age === numSupersteps)
-    sc.stop()
   }
 
   test("halting by message silence") {
-    val sc = new SparkContext("local", "test")
+    sc = new SparkContext("local", "test")
     val verts = sc.parallelize(Array("a", "b", "c", "d").map(id => (id, new TestVertex(false, 0))))
     val msgs = sc.parallelize(Array("a" -> new TestMessage("a")))
     val numSupersteps = 5
@@ -48,6 +54,5 @@ class BagelSuite extends FunSuite with Assertions {
       }
     for ((id, vert) <- result.collect)
       assert(vert.age === numSupersteps)
-    sc.stop()
   }
 }
-- 
cgit v1.2.3