diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java index c6915d859fdc..b4e86dd3e96f 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedDartReportApiTest.java @@ -74,7 +74,6 @@ import java.io.IOException; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.UUID; @@ -229,7 +228,7 @@ public void test_getQueryReport_forCompletedDartQuery() Assertions.assertEquals("10", result); // Now fetch the report using the SQL query ID - final GetQueryReportResponse reportResponse = msqApis.getDartQueryReport(sqlQueryId, broker1); + final GetQueryReportResponse reportResponse = waitForCompletedReports(sqlQueryId, broker1).get(0); // Verify the report response Assertions.assertNotNull(reportResponse, "Report response should not be null"); @@ -334,12 +333,12 @@ public void test_getQueryReport_fromBothBrokers() Assertions.assertEquals(1, sqlClients1.getAllClients().size(), "Broker1 should have 1 client (broker2)"); Assertions.assertEquals(1, sqlClients2.getAllClients().size(), "Broker2 should have 1 client (broker1)"); - // Fetch the report from both brokers, to verify cross-broker lookup is working - final GetQueryReportResponse reportFromBroker1 = msqApis.getDartQueryReport(sqlQueryId, broker1); - final GetQueryReportResponse reportFromBroker2 = msqApis.getDartQueryReport(sqlQueryId, broker2); + // Wait for the completed report to be available from both brokers. The SQL result can be returned + // before the controller is deregistered and its completed report is published. + final List completedReports = waitForCompletedReports(sqlQueryId, broker1, broker2); // Verify the report content - for (GetQueryReportResponse report : Arrays.asList(reportFromBroker1, reportFromBroker2)) { + for (GetQueryReportResponse report : completedReports) { Assertions.assertNotNull(report); final DartQueryInfo queryInfo = (DartQueryInfo) report.getQueryInfo(); Assertions.assertEquals(sqlQueryId, queryInfo.getSqlQueryId()); @@ -606,6 +605,48 @@ private GetQueryReportResponse waitForReport(String sqlQueryId) throw new ISE("Timed out after[%,d] ms waiting for query to be in RUNNING state", timeout); } + /** + * Polls the report API on the specified brokers until completed reports are available from all of them. + */ + private List waitForCompletedReports( + final String sqlQueryId, + final EmbeddedBroker... targetBrokers + ) + { + final long timeout = 30_000; + final long deadline = System.currentTimeMillis() + timeout; + + while (System.currentTimeMillis() < deadline) { + final List reports = new ArrayList<>(targetBrokers.length); + boolean allReportsCompleted = true; + + for (final EmbeddedBroker targetBroker : targetBrokers) { + final GetQueryReportResponse report = msqApis.getDartQueryReport(sqlQueryId, targetBroker); + if (report == null + || !(report.getQueryInfo() instanceof DartQueryInfo queryInfo) + || queryInfo.getDurationMs() == null) { + allReportsCompleted = false; + break; + } + reports.add(report); + } + + if (allReportsCompleted) { + return reports; + } + + try { + Thread.sleep(100); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + throw new ISE("Timed out after[%,d] ms waiting for completed reports", timeout); + } + /** * Gets running queries from {@link #broker1} using the provided HTTP client for authentication. */