Add tests for running queries JMX stat

presto-main/src/main/java/com/facebook/presto/execution/QueryManager.java

@@ -37,4 +37,6 @@ Duration waitForStateChange(QueryId queryId, QueryState currentState, Duration m
     void cancelQuery(QueryId queryId);
 
     void cancelStage(StageId stageId);
+
+    SqlQueryManagerStats getStats();
 }

presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java

@@ -369,6 +369,7 @@ public void cancelStage(StageId stageId)
         }
     }
 
+    @Override
     @Managed
     @Flatten
     public SqlQueryManagerStats getStats()

presto-main/src/main/java/com/facebook/presto/testing/MockQueryManager.java

@@ -18,6 +18,7 @@
 import com.facebook.presto.execution.QueryInfo;
 import com.facebook.presto.execution.QueryManager;
 import com.facebook.presto.execution.QueryState;
+import com.facebook.presto.execution.SqlQueryManagerStats;
 import com.facebook.presto.execution.StageId;
 import io.airlift.units.Duration;
 
@@ -75,4 +76,10 @@ public void cancelStage(StageId stageId)
     {
         throw new UnsupportedOperationException();
     }
+
+    @Override
+    public SqlQueryManagerStats getStats()
+    {
+        throw new UnsupportedOperationException();
+    }
 }

presto-tests/src/test/java/com/facebook/presto/execution/TestQueues.java

@@ -28,6 +28,7 @@
 import static com.facebook.presto.execution.QueryState.RUNNING;
 import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static org.testng.Assert.assertEquals;
 
 public class TestQueues
 {
@@ -61,34 +62,47 @@ private void testBasic(boolean resourceGroups)
         Map<String, String> properties = builder.build();
 
         try (DistributedQueryRunner queryRunner = createQueryRunner(properties)) {
+            QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
+
             // submit first "dashboard" query
             QueryId firstDashboardQuery = createQuery(queryRunner, newDashboardSession(), LONG_LASTING_QUERY);
 
             // wait for the first "dashboard" query to start
             waitForQueryState(queryRunner, firstDashboardQuery, RUNNING);
 
+            assertEquals(queryManager.getStats().getRunningQueries(), 1);
+
             // submit second "dashboard" query
             QueryId secondDashboardQuery = createQuery(queryRunner, newDashboardSession(), LONG_LASTING_QUERY);
 
             // wait for the second "dashboard" query to be queued ("dashboard.${USER}" queue strategy only allows one "dashboard" query to be accepted for execution)
             waitForQueryState(queryRunner, secondDashboardQuery, QUEUED);
 
+            assertEquals(queryManager.getStats().getRunningQueries(), 1);
+
             // submit first non "dashboard" query
             QueryId firstNonDashboardQuery = createQuery(queryRunner, newSession(), LONG_LASTING_QUERY);
 
             // wait for the first non "dashboard" query to start
             waitForQueryState(queryRunner, firstNonDashboardQuery, RUNNING);
 
+            assertEquals(queryManager.getStats().getRunningQueries(), 2);
+
             // submit second non "dashboard" query
             QueryId secondNonDashboardQuery = createQuery(queryRunner, newSession(), LONG_LASTING_QUERY);
 
             // wait for the second non "dashboard" query to start
             waitForQueryState(queryRunner, secondNonDashboardQuery, RUNNING);
 
+            assertEquals(queryManager.getStats().getRunningQueries(), 3);
+
             // cancel first "dashboard" query, second "dashboard" query and second non "dashboard" query should start running
             cancelQuery(queryRunner, firstDashboardQuery);
             waitForQueryState(queryRunner, firstDashboardQuery, FAILED);
             waitForQueryState(queryRunner, secondDashboardQuery, RUNNING);
+
+            assertEquals(queryManager.getStats().getRunningQueries(), 3);
+            assertEquals(queryManager.getStats().getCompletedQueries().getTotalCount(), 1);
         }
     }