Add/replace missing waitUntilEmpty() calls to listener bus

andrewor14 · andrewor14 · commit b4b6100b7992 · 2014-04-23T16:04:49.000-07:00
diff --git a/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
@@ -50,9 +50,9 @@ class SparkListenerSuite extends FunSuite with LocalSparkContext with ShouldMatc
     (1 to 5).foreach { _ => bus.post(SparkListenerJobEnd(0, JobSucceeded)) }
     assert(counter.count === 0)
 
-    // Starting listener bus should flush all buffered events (asynchronously, hence the sleep)
+    // Starting listener bus should flush all buffered events
     bus.start()
-    Thread.sleep(1000)
+    assert(bus.waitUntilEmpty(WAIT_TIMEOUT_MILLIS))
     assert(counter.count === 5)
 
     // After listener bus has stopped, posting events should not increment counter
@@ -177,6 +177,7 @@ class SparkListenerSuite extends FunSuite with LocalSparkContext with ShouldMatc
     listener.stageInfos.clear()
 
     rdd3.count()
+    assert(sc.listenerBus.waitUntilEmpty(WAIT_TIMEOUT_MILLIS))
     listener.stageInfos.size should be {2} // Shuffle map stage + result stage
     val stageInfo3 = listener.stageInfos.keys.find(_.stageId == 2).get
     stageInfo3.rddInfos.size should be {2} // ShuffledRDD, MapPartitionsRDD
