diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
index 03bcee94a2b915718fdfb4591e2ba0e8d18c55c8..c15aac775096c7b67785c0338c7d0ff3a4ed82a7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
@@ -316,10 +316,12 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
   test("SPARK-11126: no memory leak when running non SQL jobs") {
     val previousStageNumber = sqlContext.listener.stageIdToStageMetrics.size
     sqlContext.sparkContext.parallelize(1 to 10).foreach(i => ())
+    sqlContext.sparkContext.listenerBus.waitUntilEmpty(10000)
     // listener should ignore the non SQL stage
     assert(sqlContext.listener.stageIdToStageMetrics.size == previousStageNumber)
 
     sqlContext.sparkContext.parallelize(1 to 10).toDF().foreach(i => ())
+    sqlContext.sparkContext.listenerBus.waitUntilEmpty(10000)
     // listener should save the SQL stage
     assert(sqlContext.listener.stageIdToStageMetrics.size == previousStageNumber + 1)
   }