fix tests

Author: smalyshev

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeService.java

@@ -185,6 +185,7 @@ public void finishSessionEarly(String sessionId, ActionListener<Void> listener)
         ExchangeSourceHandler exchangeSource = removeExchangeSourceHandler(sessionId);
         if (exchangeSource != null) {
             exchangeSource.finishEarly(false, listener);
+            exchangeSource.onFinishEarly();
         } else {
             listener.onResponse(null);
         }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSinkHandler.java

@@ -102,7 +102,7 @@ public void fetchPageAsync(boolean sourceFinished, ActionListener<ExchangeRespon
             buffer.finish(true);
             var subSource = source.get();
             if (subSource != null) {
-                subSource.finishEarly(true, ActionListener.noop());
+                subSource.finishEarly(false, ActionListener.noop());
             }
         }
         listeners.add(listener);

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSourceHandler.java

@@ -317,9 +317,6 @@ public Releasable addEmptySink() {
      * @param drainingPages whether to discard pages already fetched in the exchange
      */
     public void finishEarly(boolean drainingPages, ActionListener<Void> listener) {
-        if (finishEarlyHandler != null) {
-            finishEarlyHandler.run();
-        }
         buffer.finish(drainingPages);
         if (remoteSinks.isEmpty()) {
             listener.onResponse(null);
@@ -332,6 +329,15 @@ public void finishEarly(boolean drainingPages, ActionListener<Void> listener) {
         }
     }
 
+    /**
+     * Calls the handler signifying that the request has been completed prematurely.
+     */
+    public void onFinishEarly() {
+        if (finishEarlyHandler != null) {
+            finishEarlyHandler.run();
+        }
+    }
+
     private static class PendingInstances {
         private final AtomicInteger instances = new AtomicInteger();
         private final SubscribableListener<Void> completion = new SubscribableListener<>();

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponseTests.java

@@ -518,30 +518,65 @@ static EsqlQueryResponse fromXContent(XContentParser parser) {
         }
     }
 
+    public static int clusterDetailsSize(int numClusters) {
+        /* Example:
+        "_clusters" : {
+        "total" : 2,
+        "successful" : 2,
+        "running" : 0,
+        "skipped" : 0,
+        "partial" : 0,
+        "failed" : 0,
+        "is_partial" : false,
+        "details" : {
+        "(local)" : {
+        "status" : "successful",
+        "indices" : "logs-1",
+        "took" : 4444,
+        "_shards" : {
+          "total" : 10,
+          "successful" : 10,
+          "skipped" : 3,
+          "failed" : 0
+        }
+        },
+        "remote1" : {
+        "status" : "successful",
+        "indices" : "remote1:logs-1",
+        "took" : 4999,
+        "_shards" : {
+          "total" : 12,
+          "successful" : 12,
+          "skipped" : 5,
+          "failed" : 0
+        }
+        }
+         */
+        return numClusters * 4 + 7;
+    }
+
     public void testChunkResponseSizeColumnar() {
-        int sizeClusterDetails = 14;
         try (EsqlQueryResponse resp = randomResponse(true, null)) {
             int columnCount = resp.pages().get(0).getBlockCount();
             int bodySize = resp.pages().stream().mapToInt(p -> p.getPositionCount() * p.getBlockCount()).sum() + columnCount * 2;
-            assertChunkCount(resp, r -> 5 + sizeClusterDetails + bodySize);
+            assertChunkCount(resp, r -> 5 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize);
         }
 
         try (EsqlQueryResponse resp = randomResponseAsync(true, null, true)) {
             int columnCount = resp.pages().get(0).getBlockCount();
             int bodySize = resp.pages().stream().mapToInt(p -> p.getPositionCount() * p.getBlockCount()).sum() + columnCount * 2;
-            assertChunkCount(resp, r -> 7 + sizeClusterDetails + bodySize); // is_running
+            assertChunkCount(resp, r -> 7 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize); // is_running
         }
     }
 
     public void testChunkResponseSizeRows() {
-        int sizeClusterDetails = 14;
         try (EsqlQueryResponse resp = randomResponse(false, null)) {
             int bodySize = resp.pages().stream().mapToInt(Page::getPositionCount).sum();
-            assertChunkCount(resp, r -> 5 + sizeClusterDetails + bodySize);
+            assertChunkCount(resp, r -> 5 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize);
         }
         try (EsqlQueryResponse resp = randomResponseAsync(false, null, true)) {
             int bodySize = resp.pages().stream().mapToInt(Page::getPositionCount).sum();
-            assertChunkCount(resp, r -> 7 + sizeClusterDetails + bodySize);
+            assertChunkCount(resp, r -> 7 + clusterDetailsSize(resp.getExecutionInfo().clusterInfo.size()) + bodySize);
         }
     }
 

x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1MissingIndicesIT.java

@@ -79,6 +79,7 @@ record ExpectedCluster(String clusterAlias, String indexExpression, String statu
     void assertExpectedClustersForMissingIndicesTests(Map<String, Object> responseMap, List<ExpectedCluster> expected) {
         Map<String, ?> clusters = (Map<String, ?>) responseMap.get("_clusters");
         assertThat((int) responseMap.get("took"), greaterThan(0));
+        assertThat((boolean) clusters.get("is_partial"), is(false));
 
         Map<String, ?> detailsMap = (Map<String, ?>) clusters.get("details");
         assertThat(detailsMap.size(), is(expected.size()));

x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS1UnavailableRemotesIT.java

@@ -101,6 +101,7 @@ private void clusterShutDownWithRandomSkipUnavailable() throws Exception {
         assertThat((int) clusters.get("skipped"), is(0));
         assertThat((int) clusters.get("partial"), is(0));
         assertThat((int) clusters.get("failed"), is(0));
+        assertThat((boolean) clusters.get("is_partial"), is(false));
 
         assertThat(clusterDetails.size(), is(2));
         assertThat((int) localClusterDetails.get("took"), greaterThan(0));

x-pack/plugin/security/qa/multi-cluster/src/javaRestTest/java/org/elasticsearch/xpack/remotecluster/CrossClusterEsqlRCS2UnavailableRemotesIT.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.test.cluster.util.resource.Resource;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
+import org.hamcrest.CoreMatchers;
 import org.junit.Before;
 import org.junit.ClassRule;
 import org.junit.rules.RuleChain;
@@ -119,6 +120,7 @@ private void clusterShutDownWithRandomSkipUnavailable() throws Exception {
         assertThat((int) clusters.get("skipped"), is(0));
         assertThat((int) clusters.get("partial"), is(0));
         assertThat((int) clusters.get("failed"), is(0));
+        assertThat((boolean) clusters.get("is_partial"), is(false));
 
         assertThat(clusterDetails.size(), is(2));
         assertThat((int) localClusterDetails.get("took"), greaterThan(0));