Skip to content

Commit

Permalink
ESQL: Run async tests more carefully (elastic#104330)
Browse files Browse the repository at this point in the history 
The ESQL async tests run the ESQL yaml tests two extra time - once under
the async endpoint with the `wait_for_completion_timeout` set to a long
time and *again* with `wait_for_completion_timeout` set to a short time,
expecting to receive an `id` for the query.

That second way is tricky! Even with a `0ms` timeout sometimes the
request will complete. That's great, but the tests didn't realize that
was possible. And it's tricky to get the warnings and `catch` sections
working properly with that. This reworks how we run these commands,
breaking apart the way we run a single API and running it as two, taking
into account that the "start the query" request could also complete the
query.

Closes elastic#104294
  • Loading branch information
@nik9000
nik9000 authored Jan 12, 2024
1 parent d5ae347 commit 50ac280
Show file tree
Hide file tree
Showing 4 changed files with 133 additions and 79 deletions.
8 changes: 4 additions & 4 deletions ...ml-rest-runner/src/main/java/org/elasticsearch/test/rest/yaml/section/ApiCallSection.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,10 @@ public ApiCallSection(String api) {
this.api = api;
}

public String getApi() {
return api;
}

public ApiCallSection copyWithNewApi(String api) {
ApiCallSection copy = new ApiCallSection(api);
for (var e : params.entrySet()) {
Expand All @@ -45,10 +49,6 @@ public ApiCallSection copyWithNewApi(String api) {
return copy;
}

public String getApi() {
return api;
}

public Map<String, String> getParams() {
// make sure we never modify the parameters once returned
return unmodifiableMap(params);
Expand Down
76 changes: 43 additions & 33 deletions test/yaml-rest-runner/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,6 @@ public XContentLocation getLocation() {

@Override
public void execute(ClientYamlTestExecutionContext executionContext) throws IOException {

if ("param".equals(catchParam)) {
// client should throw validation error before sending request
// lets just return without doing anything as we don't have any client to test here
Expand All @@ -359,17 +358,7 @@ public void execute(ClientYamlTestExecutionContext executionContext) throws IOEx
apiCallSection.getHeaders(),
apiCallSection.getNodeSelector()
);
if (Strings.hasLength(catchParam)) {
String catchStatusCode;
if (CATCHES.containsKey(catchParam)) {
catchStatusCode = CATCHES.get(catchParam).v1();
} else if (catchParam.startsWith("/") && catchParam.endsWith("/")) {
catchStatusCode = "4xx|5xx";
} else {
throw new UnsupportedOperationException("catch value [" + catchParam + "] not supported");
}
fail(formatStatusCodeMessage(response, catchStatusCode));
}
failIfHasCatch(response);
final String testPath = executionContext.getClientYamlTestCandidate() != null
? executionContext.getClientYamlTestCandidate().getTestPath()
: null;
Expand All @@ -393,27 +382,23 @@ public void execute(ClientYamlTestExecutionContext executionContext) throws IOEx
}
checkWarningHeaders(response.getWarningHeaders(), testPath);
} catch (ClientYamlTestResponseException e) {
ClientYamlTestResponse restTestResponse = e.getRestTestResponse();
if (Strings.hasLength(catchParam) == false) {
fail(formatStatusCodeMessage(restTestResponse, "2xx"));
} else if (CATCHES.containsKey(catchParam)) {
assertStatusCode(restTestResponse);
} else if (catchParam.length() > 2 && catchParam.startsWith("/") && catchParam.endsWith("/")) {
// the text of the error message matches regular expression
assertThat(
formatStatusCodeMessage(restTestResponse, "4xx|5xx"),
e.getResponseException().getResponse().getStatusLine().getStatusCode(),
greaterThanOrEqualTo(400)
);
Object error = executionContext.response("error");
assertThat("error was expected in the response", error, notNullValue());
// remove delimiters from regex
String regex = catchParam.substring(1, catchParam.length() - 1);
assertThat("the error message was expected to match the provided regex but didn't", error.toString(), matches(regex));
} else {
throw new UnsupportedOperationException("catch value [" + catchParam + "] not supported");
}
checkResponseException(e, executionContext);
}
}

public void failIfHasCatch(ClientYamlTestResponse response) {
if (Strings.hasLength(catchParam) == false) {
return;
}
String catchStatusCode;
if (CATCHES.containsKey(catchParam)) {
catchStatusCode = CATCHES.get(catchParam).v1();
} else if (catchParam.startsWith("/") && catchParam.endsWith("/")) {
catchStatusCode = "4xx|5xx";
} else {
throw new UnsupportedOperationException("catch value [" + catchParam + "] not supported");
}
fail(formatStatusCodeMessage(response, catchStatusCode));
}

void checkElasticProductHeader(final List<String> productHeaders) {
Expand Down Expand Up @@ -448,7 +433,7 @@ void checkWarningHeaders(final List<String> warningHeaders) {
/**
* Check that the response contains only the warning headers that we expect.
*/
void checkWarningHeaders(final List<String> warningHeaders, String testPath) {
public void checkWarningHeaders(final List<String> warningHeaders, String testPath) {
final List<String> unexpected = new ArrayList<>();
final List<String> unmatched = new ArrayList<>();
final List<String> missing = new ArrayList<>();
Expand Down Expand Up @@ -536,6 +521,31 @@ void checkWarningHeaders(final List<String> warningHeaders, String testPath) {
}
}

public void checkResponseException(ClientYamlTestResponseException e, ClientYamlTestExecutionContext executionContext)
throws IOException {

ClientYamlTestResponse restTestResponse = e.getRestTestResponse();
if (Strings.hasLength(catchParam) == false) {
fail(formatStatusCodeMessage(restTestResponse, "2xx"));
} else if (CATCHES.containsKey(catchParam)) {
assertStatusCode(restTestResponse);
} else if (catchParam.length() > 2 && catchParam.startsWith("/") && catchParam.endsWith("/")) {
// the text of the error message matches regular expression
assertThat(
formatStatusCodeMessage(restTestResponse, "4xx|5xx"),
e.getResponseException().getResponse().getStatusLine().getStatusCode(),
greaterThanOrEqualTo(400)
);
Object error = executionContext.response("error");
assertThat("error was expected in the response", error, notNullValue());
// remove delimiters from regex
String regex = catchParam.substring(1, catchParam.length() - 1);
assertThat("the error message was expected to match the provided regex but didn't", error.toString(), matches(regex));
} else {
throw new UnsupportedOperationException("catch value [" + catchParam + "] not supported");
}
}

private static void appendBadHeaders(final StringBuilder sb, final List<String> headers, final String message) {
if (headers.isEmpty() == false) {
sb.append(message).append(" [\n");
Expand Down
16 changes: 6 additions & 10 deletions .../yamlRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/EsqlClientYamlAsyncIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Stream;

/**
* Run the ESQL yaml tests against the async esql endpoint with a 30 minute {@code wait_until_completion_timeout}.
Expand All @@ -40,11 +39,11 @@ public static Iterable<Object[]> parameters() throws Exception {
body.put("wait_for_completion_timeout", "30m");
}
doSection.setApiCallSection(copy);
return Stream.of(doSection);
return doSection;
});
}

public static Iterable<Object[]> parameters(Function<DoSection, Stream<ExecutableSection>> modify) throws Exception {
public static Iterable<Object[]> parameters(Function<DoSection, ExecutableSection> modify) throws Exception {
List<Object[]> result = new ArrayList<>();
for (Object[] orig : ESClientYamlSuiteTestCase.createParameters()) {
assert orig.length == 1;
Expand All @@ -54,7 +53,7 @@ public static Iterable<Object[]> parameters(Function<DoSection, Stream<Executabl
candidate.getTestSection().getLocation(),
candidate.getTestSection().getName(),
candidate.getTestSection().getSkipSection(),
candidate.getTestSection().getExecutableSections().stream().flatMap(e -> modifyExecutableSection(e, modify)).toList()
candidate.getTestSection().getExecutableSections().stream().map(e -> modifyExecutableSection(e, modify)).toList()
);
result.add(new Object[] { new ClientYamlTestCandidate(candidate.getRestTestSuite(), modified) });
} catch (IllegalArgumentException e) {
Expand All @@ -64,12 +63,9 @@ public static Iterable<Object[]> parameters(Function<DoSection, Stream<Executabl
return result;
}

private static Stream<ExecutableSection> modifyExecutableSection(
ExecutableSection e,
Function<DoSection, Stream<ExecutableSection>> modify
) {
private static ExecutableSection modifyExecutableSection(ExecutableSection e, Function<DoSection, ExecutableSection> modify) {
if (false == (e instanceof DoSection)) {
return Stream.of(e);
return e;
}
DoSection doSection = (DoSection) e;
String api = doSection.getApiCallSection().getApi();
Expand All @@ -78,7 +74,7 @@ private static Stream<ExecutableSection> modifyExecutableSection(
case "esql.async_query", "esql.async_query_get" -> throw new IllegalArgumentException(
"The esql yaml tests can't contain async_query or async_query_get because we modify them on the fly and *add* those."
);
default -> Stream.of(e);
default -> e;
};
}
}
112 changes: 80 additions & 32 deletions ...java/org/elasticsearch/xpack/esql/qa/single_node/EsqlClientYamlAsyncSubmitAndFetchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,54 +9,102 @@

import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;

import org.apache.lucene.tests.util.LuceneTestCase;
import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
import org.elasticsearch.test.rest.yaml.section.ApiCallSection;
import org.elasticsearch.test.rest.yaml.ClientYamlTestExecutionContext;
import org.elasticsearch.test.rest.yaml.ClientYamlTestResponse;
import org.elasticsearch.test.rest.yaml.ClientYamlTestResponseException;
import org.elasticsearch.test.rest.yaml.section.DoSection;
import org.elasticsearch.test.rest.yaml.section.ExecutableSection;
import org.elasticsearch.xcontent.XContentLocation;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

/**
* Run the ESQL yaml tests async and then fetch the results with a long wait time.
*/
@LuceneTestCase.AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/104294")
public class EsqlClientYamlAsyncSubmitAndFetchIT extends AbstractEsqlClientYamlIT {
public EsqlClientYamlAsyncSubmitAndFetchIT(final ClientYamlTestCandidate testCandidate) {
super(testCandidate);
}

@ParametersFactory
public static Iterable<Object[]> parameters() throws Exception {
return EsqlClientYamlAsyncIT.parameters(doSection -> {
ApiCallSection copy = doSection.getApiCallSection().copyWithNewApi("esql.async_query");
for (Map<String, Object> body : copy.getBodies()) {
body.put("wait_for_completion_timeout", "0ms");
body.put("keep_on_completion", true);
return EsqlClientYamlAsyncIT.parameters(DoEsqlAsync::new);
}

private static class DoEsqlAsync implements ExecutableSection {
private final DoSection original;

private DoEsqlAsync(DoSection original) {
this.original = original;
}

@Override
public XContentLocation getLocation() {
return original.getLocation();
}

@Override
public void execute(ClientYamlTestExecutionContext executionContext) throws IOException {
try {
// Start the query
List<Map<String, Object>> bodies = original.getApiCallSection().getBodies().stream().map(m -> {
Map<String, Object> body = new HashMap<>(m);
if (randomBoolean()) {
/*
* Try to force the request to go async by setting the timeout to 0.
* This doesn't *actually* force the request async - if it finishes
* super duper faster it won't get async. But that's life.
*/
body.put("wait_for_completion_timeout", "0ms");
}
return body;
}).toList();
ClientYamlTestResponse startResponse = executionContext.callApi(
"esql.async_query",
original.getApiCallSection().getParams(),
bodies,
original.getApiCallSection().getHeaders(),
original.getApiCallSection().getNodeSelector()
);

String id = (String) startResponse.evaluate("id");
boolean finishedEarly = id == null;
if (finishedEarly) {
/*
* If we finished early, make sure we don't have a "catch"
* param and expect and error. And make sure we match the
* warnings folks have asked for.
*/
original.failIfHasCatch(startResponse);
original.checkWarningHeaders(startResponse.getWarningHeaders(), testPath(executionContext));
return;
}

/*
* Ok, we didn't finish before the timeout. Fine, let's fetch the result.
*/
ClientYamlTestResponse fetchResponse = executionContext.callApi(
"esql.async_query_get",
Map.of("wait_for_completion_timeout", "30m", "id", id),
List.of(),
original.getApiCallSection().getHeaders(),
original.getApiCallSection().getNodeSelector()
);
original.failIfHasCatch(fetchResponse);
original.checkWarningHeaders(fetchResponse.getWarningHeaders(), testPath(executionContext));
} catch (ClientYamlTestResponseException e) {
original.checkResponseException(e, executionContext);
}
doSection.setApiCallSection(copy);

DoSection fetch = new DoSection(doSection.getLocation());
fetch.setApiCallSection(new ApiCallSection("esql.async_query_get"));
fetch.getApiCallSection().addParam("wait_for_completion_timeout", "30m");
fetch.getApiCallSection().addParam("id", "$body.id");

/*
* The request to start the query doesn't make warnings or errors so shift
* those to the fetch.
*/
fetch.setExpectedWarningHeaders(doSection.getExpectedWarningHeaders());
fetch.setExpectedWarningHeadersRegex(doSection.getExpectedWarningHeadersRegex());
fetch.setAllowedWarningHeaders(doSection.getAllowedWarningHeaders());
fetch.setAllowedWarningHeadersRegex(doSection.getAllowedWarningHeadersRegex());
fetch.setCatch(doSection.getCatch());
doSection.setExpectedWarningHeaders(List.of());
doSection.setExpectedWarningHeadersRegex(List.of());
doSection.setAllowedWarningHeaders(List.of());
doSection.setAllowedWarningHeadersRegex(List.of());
doSection.setCatch(null);
return Stream.of(doSection, fetch);
});
}

private String testPath(ClientYamlTestExecutionContext executionContext) {
return executionContext.getClientYamlTestCandidate() != null
? executionContext.getClientYamlTestCandidate().getTestPath()
: null;
}
}
}

0 comments on commit 50ac280

Please sign in to comment.