/*
 * 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.plugin.memory;

import com.facebook.presto.server.testing.TestingPrestoServer;
import com.facebook.presto.testing.MaterializedResult;
import com.facebook.presto.tests.DistributedQueryRunner;
import io.airlift.units.Duration;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import static com.facebook.presto.plugin.memory.MemoryQueryRunner.createQueryRunner;
import static com.facebook.presto.testing.assertions.Assert.assertEquals;
import static io.airlift.testing.Assertions.assertLessThan;
import static io.airlift.units.Duration.nanosSince;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.testng.Assert.fail;

@Test(singleThreaded = true)
public class TestMemoryWorkerCrash
{
    private DistributedQueryRunner queryRunner;

    @BeforeMethod
    public void setUp()
            throws Exception
    {
        queryRunner = createQueryRunner();
    }

    @Test
    public void tableAccessAfterWorkerCrash()
            throws Exception
    {
        queryRunner.execute("CREATE TABLE test_nation as SELECT * FROM tpch.tiny.nation");
        assertQuery("SELECT * FROM test_nation ORDER BY nationkey", "SELECT * FROM tpch.tiny.nation ORDER BY nationkey");
        closeWorker();
        assertFails("SELECT * FROM test_nation ORDER BY nationkey", "No nodes available to run query");
        queryRunner.execute("INSERT INTO test_nation SELECT * FROM tpch.tiny.nation");
        assertFails("SELECT * FROM test_nation ORDER BY nationkey", "No nodes available to run query");

        queryRunner.execute("CREATE TABLE test_region as SELECT * FROM tpch.tiny.region");
        assertQuery("SELECT * FROM test_region ORDER BY regionkey", "SELECT * FROM tpch.tiny.region ORDER BY regionkey");
    }

    private void closeWorker()
            throws Exception
    {
        int nodeCount = queryRunner.getNodeCount();
        TestingPrestoServer worker = queryRunner.getServers().stream()
                .filter(server -> !server.isCoordinator())
                .findAny()
                .orElseThrow(() -> new IllegalStateException("No worker nodes"));
        worker.close();
        waitForNodes(nodeCount - 1);
    }

    private void assertQuery(String sql, String expected)
    {
        MaterializedResult rows = queryRunner.execute(sql);
        MaterializedResult expectedRows = queryRunner.execute(expected);

        assertEquals(rows, expectedRows);
    }

    private void assertFails(String sql, String expectedMessage)
    {
        try {
            queryRunner.execute(sql);
        }
        catch (RuntimeException ex) {
            // pass
            assertEquals(ex.getMessage(), expectedMessage);
            return;
        }
        fail("Query should fail");
    }

    private void waitForNodes(int numberOfNodes)
            throws InterruptedException
    {
        long start = System.nanoTime();
        while (queryRunner.getCoordinator().refreshNodes().getActiveNodes().size() < numberOfNodes) {
            assertLessThan(nanosSince(start), new Duration(10, SECONDS));
            MILLISECONDS.sleep(10);
        }
    }
}

@@ -697,30 +697,36 @@
 test
-ing.MaterializedResult
+s.AbstractTestQueryFramework
 ;%0Aim
@@ -848,217 +848,14 @@
 ons.
-BeforeMethod;%0Aimport org.testng.annotations.Test;%0A%0Aimport static com.facebook.presto.plugin.memory.MemoryQueryRunner.createQueryRunner;%0Aimport static com.facebook.presto.testing.assertions.Assert.assertEquals;
+Test;%0A
 %0Aimp
@@ -1077,46 +1077,8 @@
 NDS;
-%0Aimport static org.testng.Assert.fail;
 %0A%0A@T
@@ -1143,113 +1143,105 @@
 ash%0A
-%7B%0A
+  
     
-private DistributedQueryRunner queryRunner;%0A%0A    @BeforeMethod%0A    public void setUp
+  extends AbstractTestQueryFramework%0A%7B%0A    protected TestMemoryWorkerCrash
 ()%0A    
+%7B%0A
         
 thro
@@ -1240,40 +1240,21 @@
     
-throws Exception%0A    %7B%0A        q
+super(MemoryQ
 uery
@@ -1259,19 +1259,18 @@
 ryRunner
- = 
+::
 createQu
@@ -1278,17 +1278,16 @@
 ryRunner
-(
 );%0A    %7D
@@ -1379,33 +1379,36 @@
 n%0A    %7B%0A        
-q
+getQ
 ueryRunner.execu
@@ -1393,32 +1393,34 @@
   getQueryRunner
+()
 .execute(%22CREATE
@@ -1451,34 +1451,24 @@
 LECT * FROM 
-tpch.tiny.
 nation%22);%0A  
@@ -1544,34 +1544,24 @@
 LECT * FROM 
-tpch.tiny.
 nation ORDER
@@ -1606,32 +1606,37 @@
 ;%0A        assert
+Query
 Fails(%22SELECT * 
@@ -1709,33 +1709,36 @@
 uery%22);%0A        
-q
+getQ
 ueryRunner.execu
@@ -1723,32 +1723,34 @@
   getQueryRunner
+()
 .execute(%22INSERT
@@ -1793,32 +1793,33 @@
 .tiny.nation%22);%0A
+%0A
         assertFa
@@ -1812,24 +1812,29 @@
       assert
+Query
 Fails(%22SELEC
@@ -1912,33 +1912,36 @@
 ery%22);%0A%0A        
-q
+getQ
 ueryRunner.execu
@@ -1926,32 +1926,34 @@
   getQueryRunner
+()
 .execute(%22CREATE
@@ -2087,34 +2087,24 @@
 LECT * FROM 
-tpch.tiny.
 region ORDER
@@ -2217,28 +2217,16 @@
 Count = 
-queryRunner.
 getNodeC
@@ -2229,16 +2229,104 @@
 odeCount
+();%0A        DistributedQueryRunner queryRunner = (DistributedQueryRunner) getQueryRunner
 ();%0A    
@@ -2649,655 +2649,173 @@
 oid 
-assertQuery(String sql, String expected)%0A    %7B%0A        MaterializedResult rows = queryRunner.execute(sql);%0A        MaterializedResult expectedRows = queryRunner.execute(expected);%0A%0A        assertEquals(rows, expectedRows);
+waitForNodes(int numberOfNodes)%0A            throws InterruptedException%0A    %7B
 %0A    
-%7D%0A%0A
     
-private void assertFails(String sql, String expectedMessage)%0A    %7B%0A        try %7B%0A            queryRunner.execute(sql);%0A        %7D%0A        catch (RuntimeException ex) %7B%0A            // pass%0A            assertEquals(ex.getMessage(), expectedMessage);%0A            return;%0A        %7D%0A        fail(%22Query should fail%22);%0A    %7D%0A%0A    private void waitForNodes(int numberOfNodes)%0A            throws InterruptedException%0A    %7B
+DistributedQueryRunner queryRunner = (DistributedQueryRunner) getQueryRunner();
 %0A   

Update TestMemoryWorkerCrash to use AbstractTestQueryFramework