Since `waitUntilEmpty` is only for test, is it possible that moving such `wait` logic to the SparkListener instances used in tests? E.g.,

```Scala
  import java.util.concurrent.{CountDownLatch, TimeUnit}

  class SaveStageAndTaskInfo extends SparkListener {
    val stageInfos = mutable.Map[StageInfo, Seq[(TaskInfo, TaskMetrics)]]()
    var taskInfoMetrics = mutable.Buffer[(TaskInfo, TaskMetrics)]()
    val latch = new CountDownLatch(1)

    override def onTaskEnd(task: SparkListenerTaskEnd) {
      val info = task.taskInfo
      val metrics = task.taskMetrics
      if (info != null && metrics != null) {
        taskInfoMetrics += ((info, metrics))
      }
    }

    override def onStageCompleted(stage: SparkListenerStageCompleted) {
      stageInfos(stage.stageInfo) = taskInfoMetrics
      taskInfoMetrics = mutable.Buffer.empty
      latch.countDown()
    }

    def waitForCompleted(timeoutMillis: Long) {
      latch.await(timeoutMillis, TimeUnit.MILLISECONDS)
    }
  }
```