Backported test for the headers fix

Author: ivancea

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java

@@ -1282,6 +1282,10 @@ public static Map<String, Object> runEsqlSync(RequestObjectBuilder requestObject
         HttpEntity entity = response.getEntity();
         Map<String, Object> json = entityToMap(entity, requestObject.contentType());
 
+        var supportsAsyncHeadersFix = hasCapabilities(adminClient(), List.of("async_query_status_headers_fix"));
+        if (supportsAsyncHeadersFix) {
+            assertNoAsyncHeaders(response);
+        }
         assertWarnings(response, assertWarnings);
 
         return json;
@@ -1312,17 +1316,18 @@ public static Map<String, Object> runEsqlAsync(
         checkKeepOnCompletion(requestObject, json, keepOnCompletion);
         String id = (String) json.get("id");
 
-        var supportsAsyncHeaders = hasCapabilities(adminClient(), List.of("async_query_status_headers"));
+        var supportsAsyncHeaders = hasCapabilities(adminClient(), List.of("async_query_status_headers_fix"));
         var supportsSuggestedCast = hasCapabilities(adminClient(), List.of("suggested_cast"));
 
+        // Check headers on initial query call
+        if (supportsAsyncHeaders) {
+            assertAsyncHeaders(response, id, (boolean) json.get("is_running"));
+        }
+
         if (id == null) {
             // no id returned from an async call, must have completed immediately and without keep_on_completion
             assertThat(requestObject.keepOnCompletion(), either(nullValue()).or(is(false)));
             assertThat((boolean) json.get("is_running"), is(false));
-            if (supportsAsyncHeaders) {
-                assertThat(response.getHeader("X-Elasticsearch-Async-Id"), nullValue());
-                assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), is("?0"));
-            }
             assertWarnings(response, assertWarnings);
             json.remove("is_running"); // remove this to not mess up later map assertions
             return Collections.unmodifiableMap(json);
@@ -1343,11 +1348,6 @@ public static Map<String, Object> runEsqlAsync(
                 assertThat(json.get("pages"), nullValue());
             }
 
-            if (supportsAsyncHeaders) {
-                assertThat(response.getHeader("X-Elasticsearch-Async-Id"), is(id));
-                assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), is(isRunning ? "?1" : "?0"));
-            }
-
             // issue a second request to "async get" the results
             Request getRequest = prepareAsyncGetRequest(id);
             getRequest.setOptions(request.getOptions());
@@ -1357,6 +1357,11 @@ public static Map<String, Object> runEsqlAsync(
 
         var result = entityToMap(entity, requestObject.contentType());
 
+        // Check headers on get call
+        if (supportsAsyncHeaders) {
+            assertAsyncHeaders(response, id, (boolean) result.get("is_running"));
+        }
+
         // assert initial contents, if any, are the same as async get contents
         if (initialColumns != null) {
             if (supportsSuggestedCast == false) {
@@ -1850,6 +1855,16 @@ private static void createIndex(String indexName, boolean lookupMode, String map
         assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
     }
 
+    private static void assertAsyncHeaders(Response response, @Nullable String asyncId, boolean isRunning) {
+        assertThat(response.getHeader("X-Elasticsearch-Async-Id"), asyncId == null ? nullValue() : equalTo(asyncId));
+        assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), isRunning ? is("?1") : is("?0"));
+    }
+
+    private static void assertNoAsyncHeaders(Response response) {
+        assertThat(response.getHeader("X-Elasticsearch-Async-Id"), nullValue());
+        assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), nullValue());
+    }
+
     public static RequestObjectBuilder requestObjectBuilder() throws IOException {
         return new RequestObjectBuilder();
     }