/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.facebook.presto.execution.resourceGroups.db;

import com.facebook.presto.resourceGroups.db.H2ResourceGroupsDao;
import com.facebook.presto.spi.QueryId;
import com.facebook.presto.tests.DistributedQueryRunner;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeGroups;
import org.testng.annotations.Test;

import static com.facebook.presto.execution.QueryState.FAILED;
import static com.facebook.presto.execution.QueryState.RUNNING;
import static com.facebook.presto.execution.TestQueryRunnerUtil.createQuery;
import static com.facebook.presto.execution.TestQueryRunnerUtil.waitForQueryState;
import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.TEST_ENVIRONMENT;
import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.TEST_ENVIRONMENT_2;
import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.adhocSession;
import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.createQueryRunner;
import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.getDao;
import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.getDbConfigUrl;

public class TestEnvironments
{
    private static final String LONG_LASTING_QUERY = "SELECT COUNT(*) FROM lineitem";
    private DistributedQueryRunner runner;

    @BeforeGroups(value = TEST_ENVIRONMENT)
    public void setUpRunnerForEnvironment1()
            throws Exception
    {
        String dbConfigUrl = getDbConfigUrl();
        H2ResourceGroupsDao dao = getDao(dbConfigUrl);
        runner = createQueryRunner(dbConfigUrl, dao, TEST_ENVIRONMENT);
    }

    @BeforeGroups(value = TEST_ENVIRONMENT_2)
    public void setUpRunnerForEnvironment2()
            throws Exception
    {
        String dbConfigUrl = getDbConfigUrl();
        H2ResourceGroupsDao dao = getDao(dbConfigUrl);
        runner = createQueryRunner(dbConfigUrl, dao, TEST_ENVIRONMENT_2);
    }

    @AfterMethod(alwaysRun = true)
    public void shutdown()
    {
        runner.close();
    }

    @Test(timeOut = 60_000, groups = TEST_ENVIRONMENT)
    public void testEnvironment1()
            throws Exception
    {
        QueryId firstQuery = createQuery(runner, adhocSession(), LONG_LASTING_QUERY);
        waitForQueryState(runner, firstQuery, RUNNING);
        QueryId secondQuery = createQuery(runner, adhocSession(), LONG_LASTING_QUERY);
        waitForQueryState(runner, secondQuery, RUNNING);
    }

    @Test(timeOut = 60_000, groups = TEST_ENVIRONMENT_2)
    public void testEnvironment2()
            throws Exception
    {
        QueryId firstQuery = createQuery(runner, adhocSession(), LONG_LASTING_QUERY);
        waitForQueryState(runner, firstQuery, RUNNING);
        QueryId secondQuery = createQuery(runner, adhocSession(), LONG_LASTING_QUERY);
        // there is no queueing in TEST_ENVIRONMENT_2, so the second query should fail right away
        waitForQueryState(runner, secondQuery, FAILED);
    }
}

@@ -780,95 +780,8 @@
 er;%0A
-import org.testng.annotations.AfterMethod;%0Aimport org.testng.annotations.BeforeGroups;%0A
 impo
@@ -1634,16 +1634,45 @@
 igUrl;%0A%0A
+@Test(singleThreaded = true)%0A
 public c
@@ -1785,94 +1785,36 @@
 m%22;%0A
-    private DistributedQueryRunner runner;%0A%0A    @BeforeGroups(value = TEST_ENVIRONMENT
+%0A    @Test(timeOut = 240_000
 )%0A  
@@ -1823,38 +1823,28 @@
 public void 
-setUpRunnerFor
+test
 Environment1
@@ -1995,612 +1995,106 @@
     
-runner = createQueryRunner(dbConfigUrl, dao, TEST_ENVIRONMENT);%0A    %7D%0A%0A    @BeforeGroups(value = TEST_ENVIRONMENT_2)%0A    public void setUpRunnerForEnvironment2()%0A            throws Exception%0A    %7B%0A        String dbConfigUrl = getDbConfigUrl();%0A        H2ResourceGroupsDao dao = getDao(dbConfigUrl);%0A        runner = createQueryRunner(dbConfigUrl, dao, TEST_ENVIRONMENT_2);%0A    %7D%0A%0A    @AfterMethod(alwaysRun = true)%0A    public void shutdown()%0A    %7B%0A        runner.close();%0A    %7D%0A%0A    @Test(timeOut = 60_000, groups = TEST_ENVIRONMENT)%0A    public void testEnvironment1()%0A            throws Exception%0A    %7B%0A
+try (DistributedQueryRunner runner = createQueryRunner(dbConfigUrl, dao, TEST_ENVIRONMENT)) %7B%0A    
     
@@ -2175,32 +2175,36 @@
 QUERY);%0A        
+    
 waitForQueryStat
@@ -2235,32 +2235,36 @@
 NNING);%0A        
+  
+  
 QueryId secondQu
@@ -2318,32 +2318,36 @@
 LASTING_QUERY);%0A
+    
         waitForQ
@@ -2383,24 +2383,34 @@
 , RUNNING);%0A
+        %7D%0A
     %7D%0A%0A    @
@@ -2428,43 +2428,15 @@
 t = 
-6
+24
 0_000
-, groups = TEST_ENVIRONMENT_2
 )%0A  
@@ -2495,32 +2495,242 @@
 Exception%0A    %7B%0A
+        String dbConfigUrl = getDbConfigUrl();%0A        H2ResourceGroupsDao dao = getDao(dbConfigUrl);%0A        try (DistributedQueryRunner runner = createQueryRunner(dbConfigUrl, dao, TEST_ENVIRONMENT_2)) %7B%0A    
         QueryId 
@@ -2799,32 +2799,36 @@
 QUERY);%0A        
+ 
+   
 waitForQueryStat
@@ -2851,32 +2851,36 @@
 uery, RUNNING);%0A
+    
         QueryId 
@@ -2958,16 +2958,20 @@
         
+    
 // there
@@ -3052,32 +3052,36 @@
 ht away%0A        
+  
+  
 waitForQueryStat
@@ -3112,16 +3112,26 @@
 AILED);%0A
+        %7D%0A
     %7D%0A%7D%0A

Create runner in test method and run as single threaded