Fixed failing test

Author: tdas

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala

@@ -207,80 +207,91 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
 
     /** Custom MemoryStream that waits for manual clock to reach a time */
     val inputData = new MemoryStream[Int](0, sqlContext) {
-      // Wait for manual clock to be 100 first time there is data
+      // getOffset should take 50 ms the first time it is called
       override def getOffset: Option[Offset] = {
         val offset = super.getOffset
         if (offset.nonEmpty) {
-          clock.waitTillTime(300)
+          clock.waitTillTime(1050)
         }
         offset
       }
 
-      // Wait for manual clock to be 300 first time there is data
+      // getBatch should take 100 ms the first time it is called
       override def getBatch(start: Option[Offset], end: Offset): DataFrame = {
-        clock.waitTillTime(600)
+        if (start.isEmpty) clock.waitTillTime(1150)
         super.getBatch(start, end)
       }
     }
 
-    // This is to make sure thatquery waits for manual clock to be 600 first time there is data
-    val mapped = inputData.toDS().as[Long].map { x =>
-      clock.waitTillTime(1100)
+    // query execution should take 350 ms the first time it is called
+    val mapped = inputData.toDS.coalesce(1).as[Long].map { x =>
+      clock.waitTillTime(1500)  // this will only wait the first time when clock < 1500
       10 / x
     }.agg(count("*")).as[Long]
 
-    case class AssertStreamExecThreadToWaitForClock()
+    case class AssertStreamExecThreadIsWaitingForTime(targetTime: Long)
       extends AssertOnQuery(q => {
         eventually(Timeout(streamingTimeout)) {
           if (q.exception.isEmpty) {
-            assert(clock.asInstanceOf[StreamManualClock].isStreamWaitingAt(clock.getTimeMillis))
+            assert(clock.isStreamWaitingFor(targetTime))
           }
         }
         if (q.exception.isDefined) {
           throw q.exception.get
         }
         true
-      }, "")
+      }, "") {
+      override def toString: String = s"AssertStreamExecThreadIsWaitingForTime($targetTime)"
+    }
+
+    case class AssertClockTime(time: Long)
+      extends AssertOnQuery(q => clock.getTimeMillis() === time, "") {
+      override def toString: String = s"AssertClockTime($time)"
+    }
 
     var lastProgressBeforeStop: StreamingQueryProgress = null
 
     testStream(mapped, OutputMode.Complete)(
-      StartStream(ProcessingTime(100), triggerClock = clock),
-      AssertStreamExecThreadToWaitForClock(),
+      StartStream(ProcessingTime(1000), triggerClock = clock),
+      AssertStreamExecThreadIsWaitingForTime(1000),
       AssertOnQuery(_.status.isDataAvailable === false),
       AssertOnQuery(_.status.isTriggerActive === false),
       AssertOnQuery(_.status.message === "Waiting for next trigger"),
       AssertOnQuery(_.recentProgress.count(_.numInputRows > 0) === 0),
 
       // Test status and progress while offset is being fetched
       AddData(inputData, 1, 2),
-      AdvanceManualClock(100), // time = 100 to start new trigger, will block on getOffset
-      AssertStreamExecThreadToWaitForClock(),
+      AdvanceManualClock(1000), // time = 1000 to start new trigger, will block on getOffset
+      AssertStreamExecThreadIsWaitingForTime(1050),
       AssertOnQuery(_.status.isDataAvailable === false),
       AssertOnQuery(_.status.isTriggerActive === true),
       AssertOnQuery(_.status.message.startsWith("Getting offsets from")),
       AssertOnQuery(_.recentProgress.count(_.numInputRows > 0) === 0),
 
       // Test status and progress while batch is being fetched
-      AdvanceManualClock(200), // time = 300 to unblock getOffset, will block on getBatch
-      AssertStreamExecThreadToWaitForClock(),
+      AdvanceManualClock(50), // time = 1050 to unblock getOffset
+      AssertClockTime(1050),
+      AssertStreamExecThreadIsWaitingForTime(1150),      // will block on getBatch that needs 1150
       AssertOnQuery(_.status.isDataAvailable === true),
       AssertOnQuery(_.status.isTriggerActive === true),
       AssertOnQuery(_.status.message === "Processing new data"),
       AssertOnQuery(_.recentProgress.count(_.numInputRows > 0) === 0),
 
       // Test status and progress while batch is being processed
-      AdvanceManualClock(300), // time = 600 to unblock getBatch, will block in Spark job
+      AdvanceManualClock(100), // time = 1150 to unblock getBatch
+      AssertClockTime(1150),
+      AssertStreamExecThreadIsWaitingForTime(1500), // will block in Spark job that needs 1500
       AssertOnQuery(_.status.isDataAvailable === true),
       AssertOnQuery(_.status.isTriggerActive === true),
       AssertOnQuery(_.status.message === "Processing new data"),
       AssertOnQuery(_.recentProgress.count(_.numInputRows > 0) === 0),
 
       // Test status and progress while batch processing has completed
-      AdvanceManualClock(500), // time = 1100 to unblock job
-      AssertOnQuery { _ => clock.getTimeMillis() === 1100 },
+      AssertOnQuery { _ => clock.getTimeMillis() === 1150 },
+      AdvanceManualClock(350), // time = 1500 to unblock job
+      AssertClockTime(1500),
       CheckAnswer(2),
-      AssertStreamExecThreadToWaitForClock(),
+      AssertStreamExecThreadIsWaitingForTime(2000),
       AssertOnQuery(_.status.isDataAvailable === true),
       AssertOnQuery(_.status.isTriggerActive === false),
       AssertOnQuery(_.status.message === "Waiting for next trigger"),
@@ -293,21 +304,21 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
         assert(progress.id === query.id)
         assert(progress.name === query.name)
         assert(progress.batchId === 0)
-        assert(progress.timestamp === "1970-01-01T00:00:00.100Z") // 100 ms in UTC
+        assert(progress.timestamp === "1970-01-01T00:00:01.000Z") // 100 ms in UTC
         assert(progress.numInputRows === 2)
-        assert(progress.processedRowsPerSecond === 2.0)
+        assert(progress.processedRowsPerSecond === 4.0)
 
-        assert(progress.durationMs.get("getOffset") === 200)
-        assert(progress.durationMs.get("getBatch") === 300)
+        assert(progress.durationMs.get("getOffset") === 50)
+        assert(progress.durationMs.get("getBatch") === 100)
         assert(progress.durationMs.get("queryPlanning") === 0)
         assert(progress.durationMs.get("walCommit") === 0)
-        assert(progress.durationMs.get("triggerExecution") === 1000)
+        assert(progress.durationMs.get("triggerExecution") === 500)
 
         assert(progress.sources.length === 1)
         assert(progress.sources(0).description contains "MemoryStream")
         assert(progress.sources(0).startOffset === null)
         assert(progress.sources(0).endOffset !== null)
-        assert(progress.sources(0).processedRowsPerSecond === 2.0)
+        assert(progress.sources(0).processedRowsPerSecond === 4.0)  // 2 rows processed in 500 ms
 
         assert(progress.stateOperators.length === 1)
         assert(progress.stateOperators(0).numRowsUpdated === 1)
@@ -317,23 +328,27 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
         true
       },
 
+      // Test whether input rate is updated after two batches
+      AssertStreamExecThreadIsWaitingForTime(2000),  // blocked waiting for next trigger time
       AddData(inputData, 1, 2),
-      AdvanceManualClock(100), // allow another trigger
-      AssertStreamExecThreadToWaitForClock(),
+      AdvanceManualClock(500), // allow another trigger
+      AssertClockTime(2000),
+      AssertStreamExecThreadIsWaitingForTime(3000),  // will block waiting for next trigger time
       CheckAnswer(4),
       AssertOnQuery(_.status.isDataAvailable === true),
       AssertOnQuery(_.status.isTriggerActive === false),
       AssertOnQuery(_.status.message === "Waiting for next trigger"),
       AssertOnQuery { query =>
         assert(query.recentProgress.last.eq(query.lastProgress))
         assert(query.lastProgress.batchId === 1)
-        assert(query.lastProgress.sources(0).inputRowsPerSecond === 1.818)
+        assert(query.lastProgress.inputRowsPerSecond === 2.0)
+        assert(query.lastProgress.sources(0).inputRowsPerSecond === 2.0)
         true
       },
 
       // Test status and progress after data is not available for a trigger
-      AdvanceManualClock(100), // allow another trigger
-      AssertStreamExecThreadToWaitForClock(),
+      AdvanceManualClock(1000), // allow another trigger
+      AssertStreamExecThreadIsWaitingForTime(4000),
       AssertOnQuery(_.status.isDataAvailable === false),
       AssertOnQuery(_.status.isTriggerActive === false),
       AssertOnQuery(_.status.message === "Waiting for next trigger"),
@@ -350,10 +365,10 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
       AssertOnQuery(_.status.message === "Stopped"),
 
       // Test status and progress after query terminated with error
-      StartStream(ProcessingTime(100), triggerClock = clock),
-      AdvanceManualClock(100), // ensure initial trigger completes before AddData
+      StartStream(ProcessingTime(1000), triggerClock = clock),
+      AdvanceManualClock(1000), // ensure initial trigger completes before AddData
       AddData(inputData, 0),
-      AdvanceManualClock(100), // allow another trigger
+      AdvanceManualClock(1000), // allow another trigger
       ExpectFailure[SparkException](),
       AssertOnQuery(_.status.isDataAvailable === false),
       AssertOnQuery(_.status.isTriggerActive === false),
@@ -678,5 +693,5 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
 
 object StreamingQuerySuite {
   // Singleton reference to clock that does not get serialized in task closures
-  var clock: ManualClock = null
+  var clock: StreamManualClock = null
 }

sql/core/src/test/scala/org/apache/spark/sql/streaming/util/StreamManualClock.scala

@@ -19,7 +19,10 @@ package org.apache.spark.sql.streaming.util
 
 import org.apache.spark.util.ManualClock
 
-/** ManualClock used for streaming tests */
+/**
+ * ManualClock used for streaming tests that allows checking whether the stream is waiting
+ * on the clock at expected times.
+ */
 class StreamManualClock(time: Long = 0L) extends ManualClock(time) with Serializable {
   private var waitStartTime: Option[Long] = None
   private var waitTargetTime: Option[Long] = None