add status assertion in new ut case

Author: ruanwenjun

externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/ExecuteStatement.scala

@@ -129,6 +129,8 @@ class ExecuteStatement(
       if (jdbcStatement != null) {
         dialect.cancelStatement(jdbcStatement)
         jdbcStatement = null
+      } else {
+        warn(s"The cancel operation $statementId might be ignore due to jdbcStatement is null.")
       }
     }
   }

externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/mysql/OperationWithEngineSuite.scala

@@ -23,6 +23,7 @@ import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.engine.jdbc.connection.ConnectionProvider
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TOperationState.{CANCELED_STATE, RUNNING_STATE}
 
 class OperationWithEngineSuite extends MySQLOperationSuite with HiveJDBCTestHelper {
 
@@ -90,13 +91,22 @@ class OperationWithEngineSuite extends MySQLOperationSuite with HiveJDBCTestHelp
         val tExecuteStatementResp = client.ExecuteStatement(tExecuteStatementReq)
         assert(tExecuteStatementResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
 
-        Thread.sleep(1000) // wait for statement to start executing
+        assertOperationStatusIn(
+          client,
+          tExecuteStatementResp.getOperationHandle,
+          Set(RUNNING_STATE),
+          5)
 
         val tCancelOperationReq = new TCancelOperationReq()
         tCancelOperationReq.setOperationHandle(tExecuteStatementResp.getOperationHandle)
         val TCancelOperationReq = client.CancelOperation(tCancelOperationReq)
         assert(TCancelOperationReq.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
-        // If the statement is not cancelled successfully, will block here until 120s
+
+        assertOperationStatusIn(
+          client,
+          tExecuteStatementResp.getOperationHandle,
+          Set(CANCELED_STATE),
+          5)
       }
     }
   }

externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/starrocks/StarRocksOperationWithEngineSuite.scala

@@ -24,6 +24,7 @@ import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.engine.jdbc.connection.ConnectionProvider
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TOperationState.{CANCELED_STATE, RUNNING_STATE}
 
 class StarRocksOperationWithEngineSuite extends StarRocksOperationSuite with HiveJDBCTestHelper {
 
@@ -90,14 +91,23 @@ class StarRocksOperationWithEngineSuite extends StarRocksOperationSuite with Hiv
         tExecuteStatementReq.setRunAsync(true)
         val tExecuteStatementResp = client.ExecuteStatement(tExecuteStatementReq)
 
-        Thread.sleep(1000) // wait for statement to start executing
+        assertOperationStatusIn(
+          client,
+          tExecuteStatementResp.getOperationHandle,
+          Set(RUNNING_STATE),
+          5)
 
         val tCancelOperationReq = new TCancelOperationReq()
         tCancelOperationReq.setOperationHandle(tExecuteStatementResp.getOperationHandle)
 
         val tFetchResultsResp = client.CancelOperation(tCancelOperationReq)
         assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
-        // If the statement is not cancelled successfully, will block here until 120s
+
+        assertOperationStatusIn(
+          client,
+          tExecuteStatementResp.getOperationHandle,
+          Set(CANCELED_STATE),
+          5)
       }
     }
   }

kyuubi-common/src/test/scala/org/apache/kyuubi/operation/HiveJDBCTestHelper.scala

@@ -124,11 +124,27 @@ trait HiveJDBCTestHelper extends JDBCTestHelper {
   }
 
   def waitForOperationToComplete(client: Iface, op: TOperationHandle): Unit = {
-    val req = new TGetOperationStatusReq(op)
-    var state = client.GetOperationStatus(req).getOperationState
-    eventually(timeout(90.seconds), interval(100.milliseconds)) {
-      state = client.GetOperationStatus(req).getOperationState
-      assert(!Set(INITIALIZED_STATE, PENDING_STATE, RUNNING_STATE).contains(state))
+    assertOperationStatusIn(
+      client,
+      op,
+      Set(
+        FINISHED_STATE,
+        CANCELED_STATE,
+        CLOSED_STATE,
+        ERROR_STATE,
+        UKNOWN_STATE,
+        TIMEDOUT_STATE),
+      90)
+  }
+
+  def assertOperationStatusIn(
+      client: Iface,
+      op: TOperationHandle,
+      status: Set[TOperationState],
+      timeoutInSeconds: Int): Unit = {
+    eventually(timeout(timeoutInSeconds.seconds), interval(100.milliseconds)) {
+      val state = client.GetOperationStatus(new TGetOperationStatusReq(op)).getOperationState
+      assert(status.contains(state))
     }
   }