[FOLLOWUP apache#647] Fix flaky test

pan3793 · May 27, 2021 · 28cd6da · 28cd6da
1 parent 8ce6d01
commit 28cd6da
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/...i-main/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerConnectionSuite.scala b/...i-main/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerConnectionSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.kyuubi.operation
 
 import org.apache.hive.service.rpc.thrift.{TExecuteStatementReq, TGetOperationStatusReq, TOperationState, TStatusCode}
+import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
 
 import org.apache.kyuubi.WithKyuubiServer
 import org.apache.kyuubi.config.KyuubiConf
@@ -41,10 +42,12 @@ class KyuubiOperationPerConnectionSuite extends WithKyuubiServer with BasicJDBCT
       executeStmtReq.setRunAsync(true)
       val executeStmtResp = client.ExecuteStatement(executeStmtReq)
 
-      val getOpStatusReq = new TGetOperationStatusReq(executeStmtResp.getOperationHandle)
-      val getOpStatusResp = client.GetOperationStatus(getOpStatusReq)
-      assert(getOpStatusResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
-      assert(getOpStatusResp.getOperationState === TOperationState.ERROR_STATE)
+      eventually(timeout(10.seconds), interval(500.milliseconds)) {
+        val getOpStatusReq = new TGetOperationStatusReq(executeStmtResp.getOperationHandle)
+        val getOpStatusResp = client.GetOperationStatus(getOpStatusReq)
+        assert(getOpStatusResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
+        assert(getOpStatusResp.getOperationState === TOperationState.ERROR_STATE)
+      }
     }
   }
 }