Fix UTs

Author: HeartSaVioR

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

@@ -790,7 +790,12 @@ class EventTimeWatermarkSuite extends StreamTest with BeforeAndAfter with Matche
 
   private def assertNumLateInputs(numLateInputs: Long): AssertOnQuery = AssertOnQuery { q =>
     q.processAllAvailable()
-    val progressWithData = q.recentProgress.lastOption.get
+    val progressWithData = q.recentProgress.filterNot { p =>
+      // filter out batches which are falling into one of types:
+      // 1) doesn't execute the batch run
+      // 2) empty input batch
+      p.inputRowsPerSecond == 0
+    }.lastOption.get
     assert(progressWithData.stateOperators(0).numLateInputs === numLateInputs)
     true
   }

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

@@ -55,6 +55,9 @@ trait StateStoreMetricsTest extends StreamTest {
         val allNumUpdatedRowsSinceLastCheck =
           progressesSinceLastCheck.map(_.stateOperators.map(_.numRowsUpdated))
 
+        val allNumLateInputsSinceLastCheck =
+          progressesSinceLastCheck.map(_.stateOperators.map(_.numLateInputs))
+
         lazy val debugString = "recent progresses:\n" +
           progressesSinceLastCheck.map(_.prettyJson).mkString("\n\n")
 
@@ -64,7 +67,7 @@ trait StateStoreMetricsTest extends StreamTest {
         val numUpdatedRows = arraySum(allNumUpdatedRowsSinceLastCheck, numStateOperators)
         assert(numUpdatedRows === updated, s"incorrect updates rows, $debugString")
 
-        val numLateInputs = recentProgress.last.stateOperators.map(_.numLateInputs)
+        val numLateInputs = arraySum(allNumLateInputsSinceLastCheck, numStateOperators)
         assert(numLateInputs === lateInputs, s"incorrect late inputs, $debugString")
 
         lastCheckedRecentProgressIndex = recentProgress.length - 1