Preserve response headers in ESQL CSV tests (elastic#99883)

This change is to retrieve response headers within the responding thread since we can't pass them between threads in a blocking fashion.
ChrisHegarty · Sep 25, 2023 · 8a05425 · 8a05425
1 parent 5597c1d
commit 8a05425
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java
@@ -384,11 +384,12 @@ protected void start(Driver driver, ActionListener<Void> driverListener) {
                     Driver.start(threadPool.executor(ESQL_THREAD_POOL_NAME), driver, between(1, 1000), driverListener);
                 }
             };
-            PlainActionFuture<Void> future = new PlainActionFuture<>();
-            runner.runToCompletion(drivers, ActionListener.releaseAfter(future, () -> Releasables.close(drivers)));
-            future.actionGet(TimeValue.timeValueSeconds(30));
-            var responseHeaders = threadPool.getThreadContext().getResponseHeaders();
-            return new ActualResults(columnNames, columnTypes, dataTypes, collectedPages, responseHeaders);
+            PlainActionFuture<ActualResults> future = new PlainActionFuture<>();
+            runner.runToCompletion(drivers, ActionListener.releaseAfter(future, () -> Releasables.close(drivers)).map(ignore -> {
+                var responseHeaders = threadPool.getThreadContext().getResponseHeaders();
+                return new ActualResults(columnNames, columnTypes, dataTypes, collectedPages, responseHeaders);
+            }));
+            return future.actionGet(TimeValue.timeValueSeconds(30));
         } finally {
             Releasables.close(() -> Releasables.close(drivers), exchangeSource::decRef);
             assertThat(bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST).getUsed(), equalTo(0L));