[iotdb] 01/01: add more test for QueryStateMachine

[xi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xingtanzjr pushed a commit to branch xingtanzjr/query_execution
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit db89c036f91815e84a4ddf9246510e72554d2eea
Author: Jinrui.Zhang <xi...@gmail.com>
AuthorDate: Wed Apr 6 13:45:39 2022 +0800

    add more test for QueryStateMachine
---
 .../org/apache/iotdb/db/mpp/execution/QueryExecution.java | 11 ++++++-----
 .../iotdb/db/mpp/execution/QueryStateMachineTest.java     | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/execution/QueryExecution.java b/server/src/main/java/org/apache/iotdb/db/mpp/execution/QueryExecution.java
index 7d49fe8e9a..bb37be005b 100644
--- a/server/src/main/java/org/apache/iotdb/db/mpp/execution/QueryExecution.java
+++ b/server/src/main/java/org/apache/iotdb/db/mpp/execution/QueryExecution.java
@@ -34,8 +34,6 @@ import org.apache.iotdb.db.mpp.sql.statement.Statement;
 import org.apache.iotdb.rpc.RpcUtils;
 import org.apache.iotdb.rpc.TSStatusCode;
 
-import com.google.common.util.concurrent.ListenableFuture;
-
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
@@ -138,13 +136,16 @@ public class QueryExecution {
     releaseResource();
   }
 
-  /** Release the resources that current QueryExecution hold. */
-  private void releaseResource() {}
+  /**
+   * Release the resources that current QueryExecution hold.
+   */
+  private void releaseResource() {
+  }
 
   /**
    * This method will be called by the request thread from client connection. This method will block
    * until one of these conditions occurs: 1. There is a batch of result 2. There is no more result
-   * 3. The query has been cancelled 4. The query is timeout This method wil l fetch the result from
+   * 3. The query has been cancelled 4. The query is timeout This method will fetch the result from
    * DataStreamManager use the virtual ResultOperator's ID (This part will be designed and
    * implemented with DataStreamManager)
    */
diff --git a/server/src/test/java/org/apache/iotdb/db/mpp/execution/QueryStateMachineTest.java b/server/src/test/java/org/apache/iotdb/db/mpp/execution/QueryStateMachineTest.java
index dde44f0d64..c183655ce5 100644
--- a/server/src/test/java/org/apache/iotdb/db/mpp/execution/QueryStateMachineTest.java
+++ b/server/src/test/java/org/apache/iotdb/db/mpp/execution/QueryStateMachineTest.java
@@ -19,6 +19,7 @@
 
 package org.apache.iotdb.db.mpp.execution;
 
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
 import org.apache.iotdb.db.mpp.common.FragmentInstanceId;
@@ -32,6 +33,8 @@ import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+
 public class QueryStateMachineTest {
 
   @Test
@@ -97,6 +100,18 @@ public class QueryStateMachineTest {
     Assert.assertEquals(stateChangeCounter.get(), 2);
   }
 
+  @Test
+  public void TestGetStateChange() throws ExecutionException, InterruptedException {
+    AtomicInteger stateChangeCounter = new AtomicInteger(0);
+    QueryStateMachine stateMachine = genQueryStateMachine();
+    ListenableFuture<QueryState> future = stateMachine.getStateChange(QueryState.QUEUED);
+    future.addListener(stateChangeCounter::getAndIncrement, directExecutor());
+    Assert.assertEquals(stateChangeCounter.get(), 0);
+    stateMachine.transitionToRunning();
+    future.get();
+    Assert.assertEquals(stateChangeCounter.get(), 1);
+  }
+
   private QueryStateMachine genQueryStateMachine() {
     return new QueryStateMachine(genQueryId(), IoTDBThreadPoolFactory.newSingleThreadExecutor("TestQueryStateMachine"));
   }