Rename StreamingTestWaiter to BatchCounter

Author: JoshRosen

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

@@ -343,7 +343,7 @@ class CheckpointSuite extends TestSuiteBase {
       withStreamingContext(new StreamingContext(conf, batchDuration)) { ssc =>
         ssc.checkpoint(checkpointDir)
         clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
-        val waiter = new StreamingTestWaiter(ssc)
+        val batchCounter = new BatchCounter(ssc)
         val fileStream = ssc.textFileStream(testDir.toString)
         // Make value 3 take a large time to process, to ensure that the driver
         // shuts down in the middle of processing the 3rd batch
@@ -374,14 +374,14 @@ class CheckpointSuite extends TestSuiteBase {
           if (i != 3) {
             // Since we want to shut down while the 3rd batch is processing
             eventually(timeout(batchDuration * 5)) {
-              assert(waiter.getNumCompletedBatches === i)
+              assert(batchCounter.getNumCompletedBatches === i)
             }
           }
         }
         clock.addToTime(batchDuration.milliseconds)
         eventually(timeout(batchDuration * 5)) {
           // Wait until all files have been recorded and all batches have started
-          assert(recordedFiles(ssc) === Seq(1, 2, 3) && waiter.getNumStartedBatches === 3)
+          assert(recordedFiles(ssc) === Seq(1, 2, 3) && batchCounter.getNumStartedBatches === 3)
         }
         // Wait for a checkpoint to be written
         val fs = new Path(checkpointDir).getFileSystem(ssc.sc.hadoopConfiguration)
@@ -390,7 +390,7 @@ class CheckpointSuite extends TestSuiteBase {
         }
         ssc.stop()
         // Check that we shut down while the third batch was being processed
-        assert(waiter.getNumCompletedBatches === 2)
+        assert(batchCounter.getNumCompletedBatches === 2)
         assert(outputStream.output.flatten === Seq(1, 3))
       }
 
@@ -413,7 +413,7 @@ class CheckpointSuite extends TestSuiteBase {
           newClock.setTime(clock.currentTime())
           newClock
         }
-        val waiter = new StreamingTestWaiter(ssc)
+        val batchCounter = new BatchCounter(ssc)
         val outputStream = ssc.graph.getOutputStreams().head.asInstanceOf[TestOutputStream[Int]]
         // Check that we remember files that were recorded before the restart
         assert(recordedFiles(ssc) === Seq(1, 2, 3))
@@ -425,7 +425,7 @@ class CheckpointSuite extends TestSuiteBase {
           writeFile(i, clock)
           clock.addToTime(batchDuration.milliseconds)
           eventually(timeout(batchDuration * 5)) {
-            assert(waiter.getNumCompletedBatches === index + 1)
+            assert(batchCounter.getNumCompletedBatches === index + 1)
           }
         }
         clock.addToTime(batchDuration.milliseconds)

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

@@ -246,7 +246,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
         val clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
         // This `setTime` call ensures that the clock is past the creation time of `existingFile`
         clock.setTime(existingFile.lastModified + 1000)
-        val waiter = new StreamingTestWaiter(ssc)
+        val batchCounter = new BatchCounter(ssc)
         val fileStream = ssc.fileStream[LongWritable, Text, TextInputFormat](
           testDir.toString, (x: Path) => true, newFilesOnly = newFilesOnly).map(_._2.toString)
         val outputBuffer = new ArrayBuffer[Seq[String]] with SynchronizedBuffer[Seq[String]]
@@ -264,7 +264,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter {
           assert(file.lastModified === clock.currentTime)
           logInfo("Created file " + file)
           eventually(timeout(batchDuration * 5)) {
-            assert(waiter.getNumCompletedBatches === i)
+            assert(batchCounter.getNumCompletedBatches === i)
           }
         }
 

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

@@ -107,27 +107,26 @@ class TestOutputStreamWithPartitions[T: ClassTag](parent: DStream[T],
 }
 
 /**
- * An object that can be used to block until certain events occur, such as batch start/completion.
- * This is much less brittle than waiting on wall-clock time. Internally, this is implemented using
- * a StreamingListener.  Constructing a new instance automatically registers a StreamingListener on
+ * An object that counts the number of started / completed batches. This is implemented using a
+ * StreamingListener. Constructing a new instance automatically registers a StreamingListener on
  * the given StreamingContext.
  */
-class StreamingTestWaiter(ssc: StreamingContext) {
+class BatchCounter(ssc: StreamingContext) {
 
-  // All access to this state should be guarded by `StreamingTestWaiter.this.synchronized`
+  // All access to this state should be guarded by `BatchCounter.this.synchronized`
   private var numCompletedBatches = 0
   private var numStartedBatches = 0
 
   private val listener = new StreamingListener {
     override def onBatchStarted(batchStarted: StreamingListenerBatchStarted): Unit =
-      StreamingTestWaiter.this.synchronized {
+      BatchCounter.this.synchronized {
         numStartedBatches += 1
-        StreamingTestWaiter.this.notifyAll()
+        BatchCounter.this.notifyAll()
       }
     override def onBatchCompleted(batchCompleted: StreamingListenerBatchCompleted): Unit =
-      StreamingTestWaiter.this.synchronized {
+      BatchCounter.this.synchronized {
         numCompletedBatches += 1
-        StreamingTestWaiter.this.notifyAll()
+        BatchCounter.this.notifyAll()
       }
   }
   ssc.addStreamingListener(listener)