Skip to content

Commit 3746d99

Browse files
smalyshevsmalyshev
committed
Fix esql tests to allow is_partial in result
1 parent 8825fe5 commit 3746d99

File tree

9 files changed

+263
-426
lines changed

9 files changed

+263
-426
lines changed

‎docs/reference/esql/esql-across-clusters.asciidoc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -287,8 +287,8 @@ which clusters have slower response times than others.
287287
<6> The shard details for the search on that cluster, including a count of shards that were
288288
skipped due to the can-match phase results. Shards are skipped when they cannot have any matching data
289289
and therefore are not included in the full ES|QL query.
290-
<7> The `is_partial` field is set to `true` if the search was interrupted before finishing using
291-
<<esql-async-query-stop-api,async stop API>>.
290+
<7> The `is_partial` field is set to `true` if the search has partial results for any reason,
291+
for example if it was interrupted before finishing using <<esql-async-query-stop-api,async stop API>>.
292292

293293

294294
The cross-cluster metadata can be used to determine whether any data came back from a cluster.

‎test/framework/src/main/java/org/elasticsearch/test/rest/ESRestTestCase.java‎

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
import org.elasticsearch.rest.RestStatus;
7777
import org.elasticsearch.test.AbstractBroadcastResponseTestCase;
7878
import org.elasticsearch.test.ESTestCase;
79+
import org.elasticsearch.test.MapMatcher;
7980
import org.elasticsearch.xcontent.ConstructingObjectParser;
8081
import org.elasticsearch.xcontent.DeprecationHandler;
8182
import org.elasticsearch.xcontent.NamedXContentRegistry;
@@ -84,6 +85,7 @@
8485
import org.elasticsearch.xcontent.XContentParserConfiguration;
8586
import org.elasticsearch.xcontent.XContentType;
8687
import org.elasticsearch.xcontent.json.JsonXContent;
88+
import org.hamcrest.Matcher;
8789
import org.hamcrest.Matchers;
8890
import org.junit.After;
8991
import org.junit.AfterClass;
@@ -133,14 +135,13 @@
133135
import static org.elasticsearch.client.RestClient.IGNORE_RESPONSE_CODES_PARAM;
134136
import static org.elasticsearch.cluster.ClusterState.VERSION_INTRODUCING_TRANSPORT_VERSIONS;
135137
import static org.elasticsearch.core.Strings.format;
138+
import static org.elasticsearch.test.ListMatcher.matchesList;
139+
import static org.elasticsearch.test.MapMatcher.assertMap;
140+
import static org.elasticsearch.test.MapMatcher.matchesMap;
136141
import static org.elasticsearch.test.rest.TestFeatureService.ALL_FEATURES;
137142
import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;
138-
import static org.hamcrest.Matchers.anyOf;
139-
import static org.hamcrest.Matchers.containsString;
140-
import static org.hamcrest.Matchers.equalTo;
141-
import static org.hamcrest.Matchers.everyItem;
142-
import static org.hamcrest.Matchers.in;
143-
import static org.hamcrest.Matchers.notNullValue;
143+
import static org.hamcrest.Matchers.*;
144+
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
144145

145146
/**
146147
* Superclass for tests that interact with an external test cluster using Elasticsearch's {@link RestClient}.
@@ -2570,4 +2571,46 @@ public static Request newXContentRequest(HttpMethod method, String endpoint, ToX
25702571
addXContentBody(request, body);
25712572
return request;
25722573
}
2574+
2575+
protected static MapMatcher getResultMatcher(boolean includeMetadata, boolean includePartial) {
2576+
MapMatcher mapMatcher = matchesMap();
2577+
if (includeMetadata) {
2578+
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
2579+
}
2580+
// Older version may not have is_partial
2581+
if (includePartial) {
2582+
mapMatcher = mapMatcher.entry("is_partial", false);
2583+
}
2584+
return mapMatcher;
2585+
}
2586+
2587+
/**
2588+
* Create empty result matcher from result, taking into account all metadata items.
2589+
*/
2590+
protected static MapMatcher getResultMatcher(Map<String, Object> result) {
2591+
return getResultMatcher(result.containsKey("took"), result.containsKey("is_partial"));
2592+
}
2593+
2594+
/**
2595+
* Match result columns and values, with default matchers for metadata.
2596+
*/
2597+
protected static void assertResultMap(Map<String, Object> result, Matcher<?> columnMatcher, Matcher<?> valuesMatcher) {
2598+
assertMap(result, getResultMatcher(result).entry("columns", columnMatcher).entry("values", valuesMatcher));
2599+
}
2600+
2601+
protected static void assertResultMap(Map<String, Object> result, Object columnMatcher, Object valuesMatcher) {
2602+
assertMap(result, getResultMatcher(result).entry("columns", columnMatcher).entry("values", valuesMatcher));
2603+
}
2604+
2605+
/**
2606+
* Match result columns and values, with default matchers for metadata.
2607+
*/
2608+
protected static void assertResultMap(
2609+
Map<String, Object> result,
2610+
MapMatcher mapMatcher,
2611+
Matcher<?> columnMatcher,
2612+
Matcher<?> valuesMatcher
2613+
) {
2614+
assertMap(result, mapMatcher.entry("columns", columnMatcher).entry("values", valuesMatcher));
2615+
}
25732616
}

‎x-pack/plugin/esql/qa/security/src/javaRestTest/java/org/elasticsearch/xpack/esql/EsqlAsyncSecurityIT.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import java.io.IOException;
2323
import java.util.Locale;
24+
import java.util.Map;
2425

2526
import static org.elasticsearch.core.TimeValue.timeValueNanos;
2627
import static org.hamcrest.Matchers.allOf;
@@ -50,8 +51,9 @@ protected Response runESQLCommand(String user, String command) throws IOExceptio
5051
}
5152

5253
@Override
53-
protected MapMatcher responseMatcher() {
54-
return super.responseMatcher().entry("is_running", equalTo(false)).entry("id", allOf(notNullValue(), instanceOf(String.class)));
54+
protected MapMatcher responseMatcher(Map<String, Object> result) {
55+
return super.responseMatcher(result).entry("is_running", equalTo(false))
56+
.entry("id", allOf(notNullValue(), instanceOf(String.class)));
5557
}
5658

5759
@Override

‎x-pack/plugin/esql/qa/security/src/javaRestTest/java/org/elasticsearch/xpack/esql/EsqlSecurityIT.java‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -165,8 +165,8 @@ public void indexDocuments() throws IOException {
165165
}
166166
}
167167

168-
protected MapMatcher responseMatcher() {
169-
return matchesMap();
168+
protected MapMatcher responseMatcher(Map<String, Object> result) {
169+
return getResultMatcher(result);
170170
}
171171

172172
public void testAllowedIndices() throws Exception {
@@ -182,10 +182,7 @@ public void testAllowedIndices() throws Exception {
182182
Response resp = runESQLCommand(user, "from index-user1 | stats sum=sum(value)");
183183
assertOK(resp);
184184
Map<String, Object> responseMap = entityAsMap(resp);
185-
MapMatcher mapMatcher = responseMatcher();
186-
if (responseMap.get("took") != null) {
187-
mapMatcher = mapMatcher.entry("took", ((Integer) responseMap.get("took")).intValue());
188-
}
185+
MapMatcher mapMatcher = responseMatcher(responseMap);
189186
MapMatcher matcher = mapMatcher.entry("columns", List.of(Map.of("name", "sum", "type", "double")))
190187
.entry("values", List.of(List.of(43.0d)));
191188
assertMap(responseMap, matcher);
@@ -195,10 +192,7 @@ public void testAllowedIndices() throws Exception {
195192
Response resp = runESQLCommand(user, "from index-user2 | stats sum=sum(value)");
196193
assertOK(resp);
197194
Map<String, Object> responseMap = entityAsMap(resp);
198-
MapMatcher mapMatcher = responseMatcher();
199-
if (responseMap.get("took") != null) {
200-
mapMatcher = mapMatcher.entry("took", ((Integer) responseMap.get("took")).intValue());
201-
}
195+
MapMatcher mapMatcher = responseMatcher(responseMap);
202196
MapMatcher matcher = mapMatcher.entry("columns", List.of(Map.of("name", "sum", "type", "double")))
203197
.entry("values", List.of(List.of(72.0d)));
204198
assertMap(responseMap, matcher);
@@ -208,10 +202,7 @@ public void testAllowedIndices() throws Exception {
208202
Response resp = runESQLCommand("metadata1_read2", "from " + index + " | stats sum=sum(value)");
209203
assertOK(resp);
210204
Map<String, Object> responseMap = entityAsMap(resp);
211-
MapMatcher mapMatcher = responseMatcher();
212-
if (responseMap.get("took") != null) {
213-
mapMatcher = mapMatcher.entry("took", ((Integer) responseMap.get("took")).intValue());
214-
}
205+
MapMatcher mapMatcher = responseMatcher(responseMap);
215206
MapMatcher matcher = mapMatcher.entry("columns", List.of(Map.of("name", "sum", "type", "double")))
216207
.entry("values", List.of(List.of(72.0d)));
217208
assertMap(responseMap, matcher);
@@ -226,7 +217,7 @@ public void testAliases() throws Exception {
226217
);
227218
assertOK(resp);
228219
Map<String, Object> responseMap = entityAsMap(resp);
229-
MapMatcher matcher = responseMatcher().entry("took", ((Integer) responseMap.get("took")).intValue())
220+
MapMatcher matcher = responseMatcher(responseMap)
230221
.entry("columns", List.of(Map.of("name", "sum", "type", "double"), Map.of("name", "index", "type", "keyword")))
231222
.entry("values", List.of(List.of(72.0d, "index-user2")));
232223
assertMap(responseMap, matcher);
@@ -238,7 +229,7 @@ public void testAliasFilter() throws Exception {
238229
Response resp = runESQLCommand("alias_user1", "from " + index + " METADATA _index" + "| KEEP _index, org, value | LIMIT 10");
239230
assertOK(resp);
240231
Map<String, Object> responseMap = entityAsMap(resp);
241-
MapMatcher matcher = responseMatcher().entry("took", ((Integer) responseMap.get("took")).intValue())
232+
MapMatcher matcher = responseMatcher(responseMap)
242233
.entry(
243234
"columns",
244235
List.of(

‎x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClustersIT.java‎

Lines changed: 24 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -39,11 +39,7 @@
3939
import static org.elasticsearch.test.MapMatcher.assertMap;
4040
import static org.elasticsearch.test.MapMatcher.matchesMap;
4141
import static org.elasticsearch.xpack.esql.ccq.Clusters.REMOTE_CLUSTER_NAME;
42-
import static org.hamcrest.Matchers.any;
43-
import static org.hamcrest.Matchers.equalTo;
44-
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
45-
import static org.hamcrest.Matchers.hasItems;
46-
import static org.hamcrest.Matchers.hasKey;
42+
import static org.hamcrest.Matchers.*;
4743

4844
@ThreadLeakFilters(filters = TestClustersThreadFilter.class)
4945
public class MultiClustersIT extends ESRestTestCase {
@@ -160,42 +156,33 @@ private Map<String, Object> runEsql(RestEsqlTestCase.RequestObjectBuilder reques
160156
}
161157
}
162158

159+
private <C, V> void assertResultMap(boolean includeCCSMetadata, Map<String, Object> result, C columns, V values, boolean remoteOnly) {
160+
MapMatcher mapMatcher = getResultMatcher(ccsMetadataAvailable(), result.containsKey("is_partial"));
161+
if (includeCCSMetadata) {
162+
mapMatcher = mapMatcher.entry("_clusters", any(Map.class));
163+
}
164+
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
165+
if (includeCCSMetadata) {
166+
assertClusterDetailsMap(result, remoteOnly);
167+
}
168+
}
169+
163170
public void testCount() throws Exception {
164171
{
165172
boolean includeCCSMetadata = includeCCSMetadata();
166173
Map<String, Object> result = run("FROM test-local-index,*:test-remote-index | STATS c = COUNT(*)", includeCCSMetadata);
167174
var columns = List.of(Map.of("name", "c", "type", "long"));
168175
var values = List.of(List.of(localDocs.size() + remoteDocs.size()));
169176

170-
MapMatcher mapMatcher = matchesMap();
171-
if (includeCCSMetadata) {
172-
mapMatcher = mapMatcher.entry("_clusters", any(Map.class));
173-
}
174-
if (ccsMetadataAvailable()) {
175-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
176-
}
177-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
178-
if (includeCCSMetadata) {
179-
assertClusterDetailsMap(result, false);
180-
}
177+
assertResultMap(includeCCSMetadata, result, columns, values, false);
181178
}
182179
{
183180
boolean includeCCSMetadata = includeCCSMetadata();
184181
Map<String, Object> result = run("FROM *:test-remote-index | STATS c = COUNT(*)", includeCCSMetadata);
185182
var columns = List.of(Map.of("name", "c", "type", "long"));
186183
var values = List.of(List.of(remoteDocs.size()));
187184

188-
MapMatcher mapMatcher = matchesMap();
189-
if (includeCCSMetadata) {
190-
mapMatcher = mapMatcher.entry("_clusters", any(Map.class));
191-
}
192-
if (ccsMetadataAvailable()) {
193-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
194-
}
195-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
196-
if (includeCCSMetadata) {
197-
assertClusterDetailsMap(result, true);
198-
}
185+
assertResultMap(includeCCSMetadata, result, columns, values, true);
199186
}
200187
}
201188

@@ -208,17 +195,7 @@ public void testUngroupedAggs() throws Exception {
208195
var values = List.of(List.of(Math.toIntExact(sum)));
209196

210197
// check all sections of map except _cluster/details
211-
MapMatcher mapMatcher = matchesMap();
212-
if (includeCCSMetadata) {
213-
mapMatcher = mapMatcher.entry("_clusters", any(Map.class));
214-
}
215-
if (ccsMetadataAvailable()) {
216-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
217-
}
218-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
219-
if (includeCCSMetadata) {
220-
assertClusterDetailsMap(result, false);
221-
}
198+
assertResultMap(includeCCSMetadata, result, columns, values, false);
222199
}
223200
{
224201
boolean includeCCSMetadata = includeCCSMetadata();
@@ -227,17 +204,7 @@ public void testUngroupedAggs() throws Exception {
227204
long sum = remoteDocs.stream().mapToLong(d -> d.data).sum();
228205
var values = List.of(List.of(Math.toIntExact(sum)));
229206

230-
MapMatcher mapMatcher = matchesMap();
231-
if (includeCCSMetadata) {
232-
mapMatcher = mapMatcher.entry("_clusters", any(Map.class));
233-
}
234-
if (ccsMetadataAvailable()) {
235-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
236-
}
237-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
238-
if (includeCCSMetadata) {
239-
assertClusterDetailsMap(result, true);
240-
}
207+
assertResultMap(includeCCSMetadata, result, columns, values, true);
241208
}
242209
{
243210
assumeTrue("requires ccs metadata", ccsMetadataAvailable());
@@ -246,23 +213,15 @@ public void testUngroupedAggs() throws Exception {
246213
long sum = remoteDocs.stream().mapToLong(d -> d.data).sum();
247214
var values = List.of(List.of(Math.toIntExact(sum)));
248215

249-
MapMatcher mapMatcher = matchesMap();
250-
assertMap(
251-
result,
252-
mapMatcher.entry("columns", columns)
253-
.entry("values", values)
254-
.entry("took", greaterThanOrEqualTo(0))
255-
.entry("_clusters", any(Map.class))
256-
);
257-
assertClusterDetailsMap(result, true);
216+
assertResultMap(true, result, columns, values, true);
258217
}
259218
}
260219

261220
private void assertClusterDetailsMap(Map<String, Object> result, boolean remoteOnly) {
262221
@SuppressWarnings("unchecked")
263222
Map<String, Object> clusters = (Map<String, Object>) result.get("_clusters");
264-
assertThat(clusters.size(), greaterThanOrEqualTo(7));
265-
assertThat(clusters.keySet(), hasItems("total", "successful", "running", "skipped", "partial", "failed", "details"));
223+
assertThat(clusters.size(), equalTo(7));
224+
assertThat(clusters.keySet(), equalTo(Set.of("total", "successful", "running", "skipped", "partial", "failed", "details")));
266225
int expectedNumClusters = remoteOnly ? 1 : 2;
267226
Set<String> expectedClusterAliases = remoteOnly ? Set.of("remote_cluster") : Set.of("remote_cluster", "(local)");
268227

@@ -272,10 +231,6 @@ private void assertClusterDetailsMap(Map<String, Object> result, boolean remoteO
272231
assertThat(clusters.get("skipped"), equalTo(0));
273232
assertThat(clusters.get("partial"), equalTo(0));
274233
assertThat(clusters.get("failed"), equalTo(0));
275-
if (result.containsKey("is_partial")) {
276-
// for some BWC tests, the is_partial key may not be present
277-
assertThat(result.get("is_partial"), equalTo(false));
278-
}
279234

280235
@SuppressWarnings("unchecked")
281236
Map<String, Object> details = (Map<String, Object>) clusters.get("details");
@@ -330,17 +285,7 @@ public void testGroupedAggs() throws Exception {
330285
.map(e -> List.of(Math.toIntExact(e.getValue()), e.getKey()))
331286
.toList();
332287

333-
MapMatcher mapMatcher = matchesMap();
334-
if (includeCCSMetadata) {
335-
mapMatcher = mapMatcher.entry("_clusters", any(Map.class));
336-
}
337-
if (ccsMetadataAvailable()) {
338-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
339-
}
340-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
341-
if (includeCCSMetadata) {
342-
assertClusterDetailsMap(result, false);
343-
}
288+
assertResultMap(includeCCSMetadata, result, columns, values, false);
344289
}
345290
{
346291
boolean includeCCSMetadata = includeCCSMetadata();
@@ -358,17 +303,7 @@ public void testGroupedAggs() throws Exception {
358303
.toList();
359304

360305
// check all sections of map except _clusters/details
361-
MapMatcher mapMatcher = matchesMap();
362-
if (includeCCSMetadata) {
363-
mapMatcher = mapMatcher.entry("_clusters", any(Map.class));
364-
}
365-
if (ccsMetadataAvailable()) {
366-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
367-
}
368-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
369-
if (includeCCSMetadata) {
370-
assertClusterDetailsMap(result, true);
371-
}
306+
assertResultMap(includeCCSMetadata, result, columns, values, true);
372307
}
373308
}
374309

@@ -383,23 +318,16 @@ public void testIndexPattern() throws Exception {
383318
Map<String, Object> result = run("FROM " + indexPattern + " | STATS c = COUNT(*)", false);
384319
var columns = List.of(Map.of("name", "c", "type", "long"));
385320
var values = List.of(List.of(localDocs.size() + remoteDocs.size()));
386-
MapMatcher mapMatcher = matchesMap();
387-
if (ccsMetadataAvailable()) {
388-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
389-
}
390-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
321+
322+
assertResultMap(false, result, columns, values, false);
391323
}
392324
{
393325
String indexPattern = randomFrom("*:test-remote-index", "*:test-remote-*", "*:test-*");
394326
Map<String, Object> result = run("FROM " + indexPattern + " | STATS c = COUNT(*)", false);
395327
var columns = List.of(Map.of("name", "c", "type", "long"));
396328
var values = List.of(List.of(remoteDocs.size()));
397329

398-
MapMatcher mapMatcher = matchesMap();
399-
if (ccsMetadataAvailable()) {
400-
mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
401-
}
402-
assertMap(result, mapMatcher.entry("columns", columns).entry("values", values));
330+
assertResultMap(false, result, columns, values, false);
403331
}
404332
}
405333

0 commit comments

Comments
 (0)