Remove sleep() in runStreamsWithPartitions(); use streaming’s Duration class.

Author: JoshRosen

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

@@ -21,8 +21,6 @@ import java.io.File
 import java.nio.charset.Charset
 
 import scala.collection.mutable.ArrayBuffer
-import scala.concurrent.duration._
-import scala.language.postfixOps
 import scala.reflect.ClassTag
 
 import com.google.common.io.Files
@@ -423,7 +421,7 @@ class CheckpointSuite extends TestSuiteBase {
     val waiter = new StreamingTestWaiter(ssc)
     ssc.start()
     // Wait for the last batch before restart to be re-processed:
-    waiter.waitForTotalBatchesCompleted(1, timeout = 10 seconds)
+    waiter.waitForTotalBatchesCompleted(1, timeout = Durations.seconds(10))
     val outputNew = advanceTimeWithRealDelay[V](ssc, nextNumBatches)
     // the first element will be re-processed data of the last batch before restart
     verifyOutput(outputNew, expectedOutput.takeRight(nextNumExpectedOutputs), useSet = true)
@@ -441,7 +439,7 @@ class CheckpointSuite extends TestSuiteBase {
     logInfo("Manual clock before advancing = " + clock.time)
     for (i <- 1 to numBatches) {
       clock.addToTime(batchDuration.milliseconds)
-      waiter.waitForTotalBatchesCompleted(i, timeout = 10 seconds)
+      waiter.waitForTotalBatchesCompleted(i, timeout = Durations.seconds(10))
     }
     logInfo("Manual clock after advancing = " + clock.time)
 

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

@@ -67,7 +67,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
       Thread.sleep(500)  // This call is to allow time for the testServer to send the data to Spark
       clock.addToTime(batchDuration.milliseconds)
     }
-    waiter.waitForTotalBatchesCompleted(input.size, timeout = 10 seconds)
+    waiter.waitForTotalBatchesCompleted(input.size, timeout = Durations.seconds(10))
     logInfo("Stopping server")
     testServer.stop()
     logInfo("Stopping context")
@@ -146,7 +146,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
       inputIterator.take(2).foreach(i => queue += ssc.sparkContext.makeRDD(Seq(i)))
       clock.addToTime(batchDuration.milliseconds)
     }
-    waiter.waitForTotalBatchesCompleted(input.size, timeout = 10 seconds)
+    waiter.waitForTotalBatchesCompleted(input.size, timeout = Durations.seconds(10))
 
     logInfo("Stopping context")
     ssc.stop()
@@ -175,12 +175,12 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
     val inputIterator = input.toIterator
     inputIterator.take(3).foreach(i => queue += ssc.sparkContext.makeRDD(Seq(i)))
     clock.addToTime(batchDuration.milliseconds)
-    waiter.waitForTotalBatchesCompleted(1, timeout = 10 seconds)
+    waiter.waitForTotalBatchesCompleted(1, timeout = Durations.seconds(10))
 
     // Enqueue the remaining items (again one by one), merged in the final batch
     inputIterator.foreach(i => queue += ssc.sparkContext.makeRDD(Seq(i)))
     clock.addToTime(batchDuration.milliseconds)
-    waiter.waitForTotalBatchesCompleted(2, timeout = 10 seconds)
+    waiter.waitForTotalBatchesCompleted(2, timeout = Durations.seconds(10))
     logInfo("Stopping context")
     ssc.stop()
 

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

@@ -22,7 +22,6 @@ import java.util.concurrent.TimeoutException
 
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.mutable.SynchronizedBuffer
-import scala.concurrent.duration.FiniteDuration
 import scala.reflect.ClassTag
 
 import org.scalatest.{BeforeAndAfter, FunSuite}
@@ -135,12 +134,12 @@ class TestOutputStreamWithPartitions[T: ClassTag](parent: DStream[T],
    */
   def waitForTotalBatchesCompleted(
       targetNumBatches: Int,
-      timeout: FiniteDuration): Unit = this.synchronized {
-    val startTime = System.nanoTime
+      timeout: Duration): Unit = this.synchronized {
+    val startTime = System.currentTimeMillis()
     def successful = getNumCompletedBatches >= targetNumBatches
-    def timedOut = (System.nanoTime - startTime) >= timeout.toNanos
+    def timedOut = (System.currentTimeMillis() - startTime) >= timeout.milliseconds
     while (!timedOut && !successful) {
-      this.wait(timeout.toMillis)
+      this.wait(timeout.milliseconds)
     }
     if (!successful && timedOut) {
       throw new TimeoutException(s"Waited for $targetNumBatches completed batches, but only" +
@@ -332,6 +331,7 @@ trait TestSuiteBase extends FunSuite with BeforeAndAfter with Logging {
     val output = outputStream.output
 
     try {
+      val waiter = new StreamingTestWaiter(ssc)
       // Start computation
       ssc.start()
 
@@ -342,7 +342,7 @@ trait TestSuiteBase extends FunSuite with BeforeAndAfter with Logging {
         for (i <- 1 to numBatches) {
           logInfo("Actually waiting for " + batchDuration)
           clock.addToTime(batchDuration.milliseconds)
-          Thread.sleep(batchDuration.milliseconds)
+          waiter.waitForTotalBatchesCompleted(i, timeout = batchDuration * 5)
         }
       } else {
         clock.addToTime(numBatches * batchDuration.milliseconds)