From b1449ea69a56239f4e0c0b68c06e6bf18a659706 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Mon, 3 Feb 2025 16:24:41 -0800
Subject: [PATCH] Wait for exchange source to complete before verifying results
 (#121603)

Closes #118732
---
 .../compute/operator/exchange/ExchangeServiceTests.java         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java
index fffeeac4e4cc2..d5c6d196a1b99 100644
--- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java
+++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java
@@ -436,8 +436,8 @@ public void testExchangeSourceContinueOnFailure() {
             exchangeSink
         );
         assertThat(actualSeqNos, equalTo(expectedSeqNos));
+        safeGet(sourceCompletionFuture);
         assertThat(completedSinks.get() + failedSinks.get(), equalTo(totalSinks.get()));
-        sourceCompletionFuture.actionGet();
         if (failedRequests.get() > 0) {
             assertThat(failedSinks.get(), greaterThan(0));
         } else {