Replace several TestWaiter methods w/ ScalaTest eventually.

Author: JoshRosen

streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala

@@ -29,6 +29,7 @@ import org.apache.hadoop.fs.{FileSystem, Path}
 import org.apache.hadoop.io.{IntWritable, Text}
 import org.apache.hadoop.mapred.TextOutputFormat
 import org.apache.hadoop.mapreduce.lib.output.{TextOutputFormat => NewTextOutputFormat}
+import org.scalatest.concurrent.Eventually._
 
 import org.apache.spark.streaming.dstream.{DStream, FileInputDStream}
 import org.apache.spark.streaming.util.ManualClock
@@ -360,19 +361,25 @@ class CheckpointSuite extends TestSuiteBase {
         outputStream.register()
         ssc.start()
 
-        clock.addToTime(batchDuration.milliseconds)
+        // Advance half a batch so that the first file is created after the StreamingContext starts
+        clock.addToTime(batchDuration.milliseconds / 2)
         // Create files and advance manual clock to process them
         for (i <- Seq(1, 2, 3)) {
           writeFile(i, clock)
           clock.addToTime(batchDuration.milliseconds)
           if (i != 3) {
             // Since we want to shut down while the 3rd batch is processing
-            waiter.waitForTotalBatchesCompleted(i, batchDuration * 5)
+            eventually(timeout(batchDuration * 5)) {
+              assert(waiter.getNumCompletedBatches === i)
+            }
           }
         }
         clock.addToTime(batchDuration.milliseconds)
-        waiter.waitForTotalBatchesStarted(3, batchDuration * 5)
         Thread.sleep(1000) // To wait for execution to actually begin
+        eventually(timeout(batchDuration * 5)) {
+          assert(waiter.getNumStartedBatches === 3)
+        }
+        assert(waiter.getNumCompletedBatches === 2)
         logInfo("Output after first start = " + outputStream.output.mkString("[", ", ", "]"))
         assert(outputStream.output.size > 0, "No files processed before restart")
         ssc.stop()
@@ -410,7 +417,9 @@ class CheckpointSuite extends TestSuiteBase {
         for ((i, index) <- Seq(7, 8, 9).zipWithIndex) {
           writeFile(i, clock)
           clock.addToTime(batchDuration.milliseconds)
-          waiter.waitForTotalBatchesCompleted(index + 1, batchDuration * 5)
+          eventually(timeout(batchDuration * 5)) {
+            assert(waiter.getNumCompletedBatches === index + 1)
+          }
         }
         clock.addToTime(batchDuration.milliseconds)
         logInfo("Output after restart = " + outputStream.output.mkString("[", ", ", "]"))

streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala

@@ -32,6 +32,7 @@ import scala.language.postfixOps
 
 import com.google.common.io.Files
 import org.scalatest.BeforeAndAfter
+import org.scalatest.concurrent.Eventually._
 
 import org.apache.spark.Logging
 import org.apache.spark.storage.StorageLevel
@@ -262,7 +263,9 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
           assert(file.setLastModified(clock.currentTime()))
           assert(file.lastModified === clock.currentTime)
           logInfo("Created file " + file)
-          waiter.waitForTotalBatchesCompleted(i, timeout = batchDuration * 5)
+          eventually(timeout(batchDuration * 5)) {
+            assert(waiter.getNumCompletedBatches === i)
+          }
         }
 
         // Verify that all the files have been read

streaming/src/test/scala/org/apache/spark/streaming/TestSuiteBase.scala

@@ -21,10 +21,11 @@ import java.io.{ObjectInputStream, IOException}
 
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.mutable.SynchronizedBuffer
-import scala.concurrent.TimeoutException
+import scala.language.implicitConversions
 import scala.reflect.ClassTag
 
 import org.scalatest.{BeforeAndAfter, FunSuite}
+import org.scalatest.time.{Span, Milliseconds => ScalaTestMilliseconds}
 
 import org.apache.spark.streaming.dstream.{DStream, InputDStream, ForEachDStream}
 import org.apache.spark.streaming.scheduler.{StreamingListenerBatchStarted, StreamingListenerBatchCompleted, StreamingListener}
@@ -138,42 +139,6 @@ class StreamingTestWaiter(ssc: StreamingContext) {
   def getNumStartedBatches: Int = this.synchronized {
     numStartedBatches
   }
-
-  /**
-   * Block until the number of completed batches reaches the given threshold.
-   */
-  def waitForTotalBatchesCompleted(
-      targetNumBatches: Int,
-      timeout: Duration): Unit = this.synchronized {
-    val startTime = System.currentTimeMillis()
-    def successful = getNumCompletedBatches >= targetNumBatches
-    def timedOut = (System.currentTimeMillis() - startTime) >= timeout.milliseconds
-    while (!timedOut && !successful) {
-      this.wait(timeout.milliseconds)
-    }
-    if (!successful && timedOut) {
-      throw new TimeoutException(s"Waited for $targetNumBatches completed batches, but only" +
-        s" $numCompletedBatches have completed after $timeout")
-    }
-  }
-
-  /**
-   * Block until the number of started batches reaches the given threshold.
-   */
-  def waitForTotalBatchesStarted(
-    targetNumBatches: Int,
-    timeout: Duration): Unit = this.synchronized {
-    val startTime = System.currentTimeMillis()
-    def successful = getNumStartedBatches >= targetNumBatches
-    def timedOut = (System.currentTimeMillis() - startTime) >= timeout.milliseconds
-    while (!timedOut && !successful) {
-      this.wait(timeout.milliseconds)
-    }
-    if (!successful && timedOut) {
-      throw new TimeoutException(s"Waited for $targetNumBatches started batches, but only" +
-        s" $numStartedBatches have started after $timeout")
-    }
-  }
 }
 
 /**
@@ -236,6 +201,14 @@ trait TestSuiteBase extends FunSuite with BeforeAndAfter with Logging {
   before(beforeFunction)
   after(afterFunction)
 
+  /**
+   * Implicit conversion which allows streaming Durations to be used with ScalaTest methods,
+   * such as `eventually`.
+   */
+  implicit def streamingDurationToScalatestSpan(duration: Duration): Span = {
+    Span(duration.milliseconds, ScalaTestMilliseconds)
+  }
+
   /**
    * Run a block of code with the given StreamingContext and automatically
    * stop the context when the block completes or when an exception is thrown.