diff --git a/presto-testing/src/main/java/io/prestosql/testing/AbstractTestQueryFramework.java b/presto-testing/src/main/java/io/prestosql/testing/AbstractTestQueryFramework.java index 12ea51e5335a..f1455ca0d319 100644 --- a/presto-testing/src/main/java/io/prestosql/testing/AbstractTestQueryFramework.java +++ b/presto-testing/src/main/java/io/prestosql/testing/AbstractTestQueryFramework.java @@ -23,6 +23,7 @@ import io.prestosql.cost.CostCalculatorWithEstimatedExchanges; import io.prestosql.cost.CostComparator; import io.prestosql.cost.TaskCountEstimator; +import io.prestosql.execution.QueryInfo; import io.prestosql.execution.QueryManagerConfig; import io.prestosql.execution.TaskManagerConfig; import io.prestosql.execution.warnings.WarningCollector; @@ -61,6 +62,7 @@ import java.util.Optional; import java.util.OptionalLong; import java.util.function.Consumer; +import java.util.function.Supplier; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; @@ -71,9 +73,11 @@ import static io.prestosql.SystemSessionProperties.JOIN_REORDERING_STRATEGY; import static io.prestosql.sql.ParsingUtil.createParsingOptions; import static io.prestosql.sql.SqlFormatter.formatSql; +import static io.prestosql.testing.assertions.Assert.assertEventually; import static io.prestosql.transaction.TransactionBuilder.transaction; import static java.lang.String.format; import static java.util.Collections.emptyList; +import static java.util.concurrent.TimeUnit.SECONDS; import static org.assertj.core.api.Assertions.assertThat; import static org.testng.Assert.assertEquals; import static org.testng.Assert.fail; @@ -446,9 +450,13 @@ protected OperatorStats searchScanFilterAndProjectOperatorStats(QueryId queryId, }) .findOnlyElement() .getId(); - return runner.getCoordinator() + + Supplier queryInfoSupplier = () -> runner.getCoordinator() .getQueryManager() - .getFullQueryInfo(queryId) + .getFullQueryInfo(queryId); + assertEventually(new Duration(5, SECONDS), () -> assertThat(queryInfoSupplier.get().isFinalQueryInfo())); + + return queryInfoSupplier.get() .getQueryStats() .getOperatorSummaries() .stream()