address comments

apache · Dec 28, 2021 · 5d12f70 · 5d12f70
1 parent 80e1404
commit 5d12f70
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 18 deletions.
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/OperationsResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/OperationsResourceSuite.scala
@@ -34,23 +34,15 @@ class OperationsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper
 
   test("get an operation event") {
     val catalogsHandleStr = getOpHandleStr(OperationType.GET_CATALOGS)
-    var response = webTarget.path(s"api/v1/operations/$catalogsHandleStr/event")
-      .request(MediaType.APPLICATION_JSON_TYPE).get()
-    val operationEvent = response.readEntity(classOf[KyuubiOperationEvent])
-    assert(200 == response.getStatus)
     checkOpState(catalogsHandleStr, FINISHED)
 
     val statementHandleStr = getOpHandleStr(OperationType.EXECUTE_STATEMENT)
-    response = webTarget.path(s"api/v1/operations/$statementHandleStr/event")
-      .request(MediaType.APPLICATION_JSON_TYPE).get()
-    val statementEvent = response.readEntity(classOf[KyuubiOperationEvent])
-    assert(200 == response.getStatus)
     checkOpState(statementHandleStr, FINISHED)
 
     // Invalid operationHandleStr
     val invalidOperationHandle =
       statementHandleStr.replaceAll("EXECUTE_STATEMENT", "GET_TYPE_INFO")
-    response = webTarget.path(s"api/v1/operations/$invalidOperationHandle/event")
+    val response = webTarget.path(s"api/v1/operations/$invalidOperationHandle/event")
       .request(MediaType.APPLICATION_JSON_TYPE).get()
     assert(404 == response.getStatus)
   }
@@ -62,18 +54,12 @@ class OperationsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper
       .request(MediaType.APPLICATION_JSON_TYPE)
       .put(Entity.entity(OpActionRequest("cancel"), MediaType.APPLICATION_JSON_TYPE))
     assert(200 == response.getStatus)
-
-    response = webTarget.path(s"api/v1/operations/$opHandleStr/event")
-      .request(MediaType.APPLICATION_JSON_TYPE).get()
-    val operationEvent = response.readEntity(classOf[KyuubiOperationEvent])
-    assert(operationEvent.state == OperationState.FINISHED.name() ||
-      operationEvent.state == OperationState.CANCELED.name())
+    checkOpState(opHandleStr, OperationState.CANCELED)
 
     response = webTarget.path(s"api/v1/operations/$opHandleStr")
       .request(MediaType.APPLICATION_JSON_TYPE)
       .put(Entity.entity(OpActionRequest("close"), MediaType.APPLICATION_JSON_TYPE))
     assert(200 == response.getStatus)
-
     response = webTarget.path(s"api/v1/operations/$opHandleStr/event")
       .request(MediaType.APPLICATION_JSON_TYPE).get()
     assert(404 == response.getStatus)
@@ -125,6 +111,7 @@ class OperationsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper
     eventually(Timeout(5.seconds)) {
       val response = webTarget.path(s"api/v1/operations/$opHandleStr/event")
         .request(MediaType.APPLICATION_JSON_TYPE).get()
+      assert(response.getStatus === 200)
       val operationEvent = response.readEntity(classOf[KyuubiOperationEvent])
       assert(operationEvent.state === state.name())
     }

diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/SessionsResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/SessionsResourceSuite.scala
@@ -239,7 +239,7 @@ class SessionsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
     operationHandle = response.readEntity(classOf[OperationHandle])
     assert(operationHandle.typ == OperationType.GET_COLUMNS)
 
-    var getFunctionsReq = GetFunctionsRequest("default", "default", "default")
+    val getFunctionsReq = GetFunctionsRequest("default", "default", "default")
     response = webTarget.path(s"$pathPrefix/operations/functions")
       .request(MediaType.APPLICATION_JSON_TYPE)
       .post(Entity.entity(getFunctionsReq, MediaType.APPLICATION_JSON_TYPE))
@@ -270,7 +270,7 @@ class SessionsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
       .request(MediaType.APPLICATION_JSON_TYPE)
       .post(Entity.entity(null, MediaType.APPLICATION_JSON_TYPE))
     assert(200 == response.getStatus)
-    var operationHandle = response.readEntity(classOf[OperationHandle])
+    val operationHandle = response.readEntity(classOf[OperationHandle])
     assert(operationHandle.typ == OperationType.GET_CATALOGS)
 
     val serializedOperationHandle = s"${operationHandle.identifier.publicId}|" +