Fix flaky test TestQueryManager.testQueryCpuLimit
We now run a query that will run for a longer time to give more time to the
SqlQueryManager thread to kill it. We also explicitly wait for the final
query state after running the query instead of waiting for the SqlQueryManager
thread implicitly.
@@ -15,11 +15,9 @@ package com.facebook.presto.tests;
import com.facebook.presto.execution.QueryInfo;
import com.facebook.presto.execution.QueryManager;
-import com.facebook.presto.execution.QueryState;
import com.facebook.presto.execution.TestingSessionContext;
import com.facebook.presto.spi.PrestoException;
import com.facebook.presto.spi.QueryId;
-import com.facebook.presto.testing.QueryRunner;
import com.google.common.collect.ImmutableMap;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
@@ -28,7 +26,11 @@ import org.testng.annotations.Test;
import java.util.Map;
import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
+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.spi.StandardErrorCode.EXCEEDED_CPU_LIMIT;
import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
import static com.facebook.presto.tests.tpch.TpchQueryRunner.createQueryRunner;
import static org.testng.Assert.assertEquals;
@@ -76,21 +78,24 @@ public class TestQueryManager
// cancel query
queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception"));
QueryInfo queryInfo = queryManager.getQueryInfo(queryId);
- assertEquals(queryInfo.getState(), QueryState.FAILED);
+ assertEquals(queryInfo.getState(), FAILED);
assertEquals(queryInfo.getErrorCode(), GENERIC_INTERNAL_ERROR.toErrorCode());
assertNotNull(queryInfo.getFailureInfo());
assertEquals(queryInfo.getFailureInfo().getMessage(), "mock exception");
}
- @Test(timeOut = 60_000, expectedExceptions = RuntimeException.class, expectedExceptionsMessageRegExp = "Exceeded CPU limit of .*")
+ @Test(timeOut = 60_000L)
public void testQueryCpuLimit()
throws Exception
{
- Map<String, String> properties = ImmutableMap.<String, String>builder()
- .put("query.max-cpu-time", "1ms")
- .build();
- try (QueryRunner queryRunner = createQueryRunner(properties)) {
- queryRunner.execute(TEST_SESSION, "SELECT COUNT(*), repeat(orderstatus, 1000) FROM orders GROUP BY 2");
+ Map<String, String> properties = ImmutableMap.of("query.max-cpu-time", "1ms");
+ try (DistributedQueryRunner queryRunner = createQueryRunner(properties)) {
+ QueryId queryId = createQuery(queryRunner, TEST_SESSION, "SELECT COUNT(*) FROM lineitem");
+ waitForQueryState(queryRunner, queryId, FAILED);
+ QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
+ QueryInfo queryInfo = queryManager.getQueryInfo(queryId);
+ assertEquals(queryInfo.getState(), FAILED);
+ assertEquals(queryInfo.getErrorCode(), EXCEEDED_CPU_LIMIT.toErrorCode());
}
}
}
