Skip to content

Commit 4951289

Browse files
smalyshevsmalyshev
authored
Refactor CrossClusterQueryWithPartialResultsIT tests (#124811) (#124842)
1 parent 33f2010 commit 4951289

File tree

1 file changed

+34
-52
lines changed

1 file changed

+34
-52
lines changed

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClusterQueryWithPartialResultsIT.java

+34-52
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,25 @@ void populateIndices() throws Exception {
6666
populateIndexWithFailingFields(REMOTE_CLUSTER_2, "fail-cluster2", remote2.failingShards);
6767
}
6868

69+
private void assertClusterPartial(EsqlQueryResponse resp, String clusterAlias, ClusterSetup cluster) {
70+
assertClusterPartial(resp, clusterAlias, cluster.okShards + cluster.failingShards, cluster.okShards);
71+
}
72+
73+
private void assertClusterPartial(EsqlQueryResponse resp, String clusterAlias, int totalShards, int okShards) {
74+
EsqlExecutionInfo.Cluster clusterInfo = resp.getExecutionInfo().getCluster(clusterAlias);
75+
assertThat(clusterInfo.getTotalShards(), equalTo(totalShards));
76+
assertThat(clusterInfo.getSuccessfulShards(), lessThanOrEqualTo(okShards));
77+
assertThat(clusterInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
78+
}
79+
80+
private void assertClusterSuccess(EsqlQueryResponse resp, String clusterAlias, int numShards) {
81+
EsqlExecutionInfo.Cluster clusterInfo = resp.getExecutionInfo().getCluster(clusterAlias);
82+
assertThat(clusterInfo.getTotalShards(), equalTo(numShards));
83+
assertThat(clusterInfo.getSuccessfulShards(), equalTo(numShards));
84+
assertThat(clusterInfo.getFailedShards(), equalTo(0));
85+
assertThat(clusterInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL));
86+
}
87+
6988
public void testPartialResults() throws Exception {
7089
populateIndices();
7190
EsqlQueryRequest request = new EsqlQueryRequest();
@@ -91,22 +110,10 @@ public void testPartialResults() throws Exception {
91110
assertTrue(returnedIds.add(id));
92111
assertThat(id, is(in(allIds)));
93112
}
94-
if (request.includeCCSMetadata()) {
95-
EsqlExecutionInfo.Cluster localInfo = resp.getExecutionInfo().getCluster(LOCAL_CLUSTER);
96-
assertThat(localInfo.getTotalShards(), equalTo(local.okShards + local.failingShards));
97-
assertThat(localInfo.getSuccessfulShards(), lessThanOrEqualTo(local.okShards));
98-
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
99-
100-
EsqlExecutionInfo.Cluster remote1Info = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_1);
101-
assertThat(remote1Info.getTotalShards(), equalTo(remote1.okShards + remote1.failingShards));
102-
assertThat(remote1Info.getSuccessfulShards(), lessThanOrEqualTo(remote1.okShards));
103-
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
104113

105-
EsqlExecutionInfo.Cluster remote2Info = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_2);
106-
assertThat(remote2Info.getTotalShards(), equalTo(remote2.okShards + remote2.failingShards));
107-
assertThat(remote2Info.getSuccessfulShards(), lessThanOrEqualTo(remote2.okShards));
108-
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
109-
}
114+
assertClusterPartial(resp, LOCAL_CLUSTER, local);
115+
assertClusterPartial(resp, REMOTE_CLUSTER_1, remote1);
116+
assertClusterPartial(resp, REMOTE_CLUSTER_2, remote2);
110117
}
111118
}
112119

@@ -128,22 +135,10 @@ public void testOneRemoteClusterPartial() throws Exception {
128135
assertTrue(returnedIds.add(id));
129136
}
130137
assertThat(returnedIds, equalTo(allIds));
131-
if (request.includeCCSMetadata()) {
132-
EsqlExecutionInfo.Cluster localInfo = resp.getExecutionInfo().getCluster(LOCAL_CLUSTER);
133-
assertThat(localInfo.getTotalShards(), equalTo(local.okShards));
134-
assertThat(localInfo.getSuccessfulShards(), equalTo(local.okShards));
135-
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL));
136138

137-
EsqlExecutionInfo.Cluster remote1Info = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_1);
138-
assertThat(remote1Info.getTotalShards(), equalTo(remote1.okShards));
139-
assertThat(remote1Info.getSuccessfulShards(), equalTo(remote1.okShards));
140-
assertThat(remote1Info.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL));
141-
142-
EsqlExecutionInfo.Cluster remote2Info = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_2);
143-
assertThat(remote2Info.getTotalShards(), equalTo(remote2.failingShards));
144-
assertThat(remote2Info.getSuccessfulShards(), equalTo(0));
145-
assertThat(remote2Info.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
146-
}
139+
assertClusterSuccess(resp, LOCAL_CLUSTER, local.okShards);
140+
assertClusterSuccess(resp, REMOTE_CLUSTER_1, remote1.okShards);
141+
assertClusterPartial(resp, REMOTE_CLUSTER_2, remote2.failingShards, 0);
147142
}
148143
}
149144

@@ -195,15 +190,10 @@ public void sendResponse(Exception exception) {
195190
assertTrue(returnedIds.add(id));
196191
}
197192
assertThat(returnedIds, equalTo(Sets.union(local.okIds, remote1.okIds)));
198-
if (request.includeCCSMetadata()) {
199-
EsqlExecutionInfo.Cluster localInfo = resp.getExecutionInfo().getCluster(LOCAL_CLUSTER);
200-
assertThat(localInfo.getTotalShards(), equalTo(localInfo.getTotalShards()));
201-
assertThat(localInfo.getSuccessfulShards(), equalTo(localInfo.getSuccessfulShards()));
202-
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL));
193+
assertClusterSuccess(resp, LOCAL_CLUSTER, local.okShards);
203194

204-
EsqlExecutionInfo.Cluster remoteInfo = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_1);
205-
assertThat(remoteInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
206-
}
195+
EsqlExecutionInfo.Cluster remoteInfo = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_1);
196+
assertThat(remoteInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
207197
}
208198
} finally {
209199
for (TransportService transportService : cluster(REMOTE_CLUSTER_1).getInstances(TransportService.class)) {
@@ -248,15 +238,10 @@ public void testFailToStartRequestOnRemoteCluster() throws Exception {
248238
assertTrue(returnedIds.add(id));
249239
}
250240
assertThat(returnedIds, equalTo(local.okIds));
251-
if (request.includeCCSMetadata()) {
252-
EsqlExecutionInfo.Cluster localInfo = resp.getExecutionInfo().getCluster(LOCAL_CLUSTER);
253-
assertThat(localInfo.getTotalShards(), equalTo(local.okShards));
254-
assertThat(localInfo.getSuccessfulShards(), equalTo(local.okShards));
255-
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL));
241+
assertClusterSuccess(resp, LOCAL_CLUSTER, local.okShards);
256242

257-
EsqlExecutionInfo.Cluster remoteInfo = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_1);
258-
assertThat(remoteInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
259-
}
243+
EsqlExecutionInfo.Cluster remoteInfo = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_1);
244+
assertThat(remoteInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
260245
}
261246
} finally {
262247
for (TransportService transportService : cluster(REMOTE_CLUSTER_1).getInstances(TransportService.class)) {
@@ -299,13 +284,10 @@ public void testFailSearchShardsOnLocalCluster() throws Exception {
299284
assertTrue(returnedIds.add(id));
300285
}
301286
assertThat(returnedIds, equalTo(remote1.okIds));
302-
if (request.includeCCSMetadata()) {
303-
EsqlExecutionInfo.Cluster localInfo = resp.getExecutionInfo().getCluster(LOCAL_CLUSTER);
304-
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
287+
EsqlExecutionInfo.Cluster localInfo = resp.getExecutionInfo().getCluster(LOCAL_CLUSTER);
288+
assertThat(localInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.PARTIAL));
305289

306-
EsqlExecutionInfo.Cluster remoteInfo = resp.getExecutionInfo().getCluster(REMOTE_CLUSTER_1);
307-
assertThat(remoteInfo.getStatus(), equalTo(EsqlExecutionInfo.Cluster.Status.SUCCESSFUL));
308-
}
290+
assertClusterSuccess(resp, REMOTE_CLUSTER_1, remote1.okShards);
309291
}
310292
} finally {
311293
for (TransportService transportService : cluster(LOCAL_CLUSTER).getInstances(TransportService.class)) {

0 commit comments

Comments
 (0)