From 28c10b54a859661c07c6e7791e1b952cf5f7b32f Mon Sep 17 00:00:00 2001
From: Cheng Pan <379377944@qq.com>
Date: Thu, 27 May 2021 11:04:24 +0800
Subject: [PATCH] [KYUUBI #648] [FOLLOWUP #647] Fix flaky test

<!--
Thanks for sending a pull request!

Here are some tips for you:
  1. If this is your first time, please read our contributor guidelines: https://kyuubi.readthedocs.io/en/latest/community/contributions.html
  2. If the PR is related to an issue in https://github.com/NetEase/kyuubi/issues, add '[KYUUBI #XXXX]' in your PR title, e.g., '[KYUUBI #XXXX] Your PR title ...'.
  3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][KYUUBI #XXXX] Your PR title ...'.
-->

### _Why are the changes needed?_
<!--
Please clarify why the changes are needed. For instance,
  1. If you add a feature, you can talk about the use case of it.
  2. If you fix a bug, you can clarify why it is a bug.
-->

### _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 #648 from pan3793/followup.

Closes #648

28cd6dae [Cheng Pan] [FOLLOWUP #647] Fix flaky test

Authored-by: Cheng Pan <379377944@qq.com>
Signed-off-by: ulysses-you <ulyssesyou18@gmail.com>
(cherry picked from commit b09c87bba12a5ff79d9bf97efd6aab172895affd)
Signed-off-by: ulysses-you <ulyssesyou18@gmail.com>
---
 .../operation/KyuubiOperationPerConnectionSuite.scala | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/kyuubi-main/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerConnectionSuite.scala b/kyuubi-main/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerConnectionSuite.scala
index 75da6246f6b..277c0c85885 100644
--- a/kyuubi-main/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerConnectionSuite.scala
+++ b/kyuubi-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)
+      }
     }
   }
 }