[KYUUBI apache#862] Revise the unit test for "add kyuubiStatementInfo into queue"

zhang1002 · yaooqinn · commit 286e3fbcc12f · 2021-07-23T23:29:31.000+08:00
### _Why are the changes needed?_ Revise the unit test for "add kyuubiStatementInfo into queue" ### _How was this patch tested?_ - [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible - [ ] Add screenshots for manual tests if appropriate - [ ] [Run test](https://kyuubi.readthedocs.io/en/latest/tools/testing.html#running-tests) locally before make a pull request Closes apache#862 from zhang1002/update_statementinfo_ut. Closes apache#862 7e5f8b5 [张宇翔] Revise the unit test for "add kyuubiStatementInfo into queue" e517cfc [张宇翔] Merge remote-tracking branch 'upstream/master' 18aebe7 [张宇翔] Merge remote-tracking branch 'upstream/master' f248bef [张宇翔] Merge remote-tracking branch 'upstream/master' 5ffb54f [张宇翔] Add kyuubi-spark-monitor module for nightly.yml Authored-by: 张宇翔 <zhang1002@126.com> Signed-off-by: Kent Yao <yao@apache.org>
diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/KyuubiStatementMonitorSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/KyuubiStatementMonitorSuite.scala
@@ -36,45 +36,29 @@ class KyuubiStatementMonitorSuite extends WithSparkSQLEngine with HiveJDBCTests
   override protected def jdbcUrl: String = getJdbcUrl
   override def withKyuubiConf: Map[String, String] = Map.empty
 
-  test("add kyuubiStatementInfo into queue and remove them by size type threshold") {
-    val sql = "select timestamp'2021-06-01'"
+  test("add kyuubiStatementInfo into queue") {
+    var baseSql = "select timestamp'2021-06-0"
     val total: Int = 7
     // Clear kyuubiStatementQueue first
     val getQueue = PrivateMethod[
       ArrayBlockingQueue[KyuubiStatementInfo]](Symbol("kyuubiStatementQueue"))()
     val kyuubiStatementQueue = KyuubiStatementMonitor.invokePrivate(getQueue)
     kyuubiStatementQueue.clear()
     withSessionHandle { (client, handle) =>
+
       for ( a <- 1 to total ) {
+        val sql = baseSql + a + "'"
         val req = new TExecuteStatementReq()
         req.setSessionHandle(handle)
         req.setStatement(sql)
         val tExecuteStatementResp = client.ExecuteStatement(req)
         val operationHandle = tExecuteStatementResp.getOperationHandle
-        waitForOperationToComplete(client, operationHandle)
-      }
 
-      var iterator = kyuubiStatementQueue.iterator()
-      while (iterator.hasNext) {
-        val kyuubiStatementInfo = iterator.next()
-        assert(kyuubiStatementInfo.statement !== null)
-        assert(kyuubiStatementInfo.statementId !== null)
-        assert(kyuubiStatementInfo.sessionId !== null)
-        assert(kyuubiStatementInfo.queryExecution !== null)
-        assert(kyuubiStatementInfo.stateToTime.size === 4)
+        val kyuubiStatementInfo = kyuubiStatementQueue.poll()
+        assert(
+          kyuubiStatementInfo.statementId === OperationHandle(operationHandle).identifier.toString)
+        assert(sql === kyuubiStatementInfo.statement)
       }
-      iterator = null
-
-      // Test for clear kyuubiStatementQueue
-      // This function is used for avoiding mem leak
-      val req = new TExecuteStatementReq()
-      req.setSessionHandle(handle)
-      req.setStatement(sql)
-      val tExecuteStatementResp = client.ExecuteStatement(req)
-      val operationHandle = tExecuteStatementResp.getOperationHandle
-      waitForOperationToComplete(client, operationHandle)
-
-      assert(kyuubiStatementQueue.size() === 1)
     }
   }
 
