Fixing some tests.

Author: afoucret

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/bulk/BulkInferenceExecutionState.java

@@ -47,13 +47,14 @@ public void markSeqNoAsPersisted(long seqNo) {
         checkpoint.markSeqNoAsPersisted(seqNo);
     }
 
-    public void persistsResponse(CheckedConsumer<InferenceAction.Response, Exception> persister) {
+    public void persistsResponses(CheckedConsumer<InferenceAction.Response, Exception> persister) {
         synchronized (checkpoint) {
             long persistedSeqNo = checkpoint.getPersistedCheckpoint();
             while (persistedSeqNo < checkpoint.getProcessedCheckpoint()) {
                 persistedSeqNo++;
                 InferenceAction.Response response = bufferedResponses.remove(persistedSeqNo);
-                if (hasFailure() == false) {
+                assert response != null || hasFailure();
+                if (hasFailure() == false && responseSent() == false) {
                     try {
                         persister.accept(response);
                     } catch (Exception e) {
@@ -90,14 +91,14 @@ public <OutputType> void maybeSendResponse(CheckedSupplier<OutputType, Exception
             if (failureCollector.hasFailure() == false) {
                 try {
                     l.onResponse(responseBuilder.get());
+                    return;
                 } catch (Exception e) {
-                    failureCollector.unwrapAndCollect(e);
+                    l.onFailure(e);
+                    return;
                 }
             }
 
-            if (failureCollector.hasFailure()) {
-                l.onFailure(failureCollector.getFailure());
-            }
+            l.onFailure(failureCollector.getFailure());
         }
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/bulk/BulkInferenceExecutor.java

@@ -51,11 +51,7 @@ public <InferenceResult extends InferenceServiceResults, OutputType> void execut
                             e -> bulkExecutionState.onInferenceException(seqNo, e)
                         ),
                         () -> {
-                            if (bulkExecutionState.responseSent()) {
-                                return;
-                            }
-
-                            bulkExecutionState.persistsResponse(outputBuilder::onInferenceResponse);
+                            bulkExecutionState.persistsResponses(outputBuilder::onInferenceResponse);
                             bulkExecutionState.maybeSendResponse(outputBuilder::buildOutput, listener);
                         }
                     )

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/bulk/BulkInferenceExecutorTests.java

@@ -67,7 +67,7 @@ public void shutdownThreadPool() {
 
     @SuppressWarnings("unchecked")
     public void testSuccessfulExecution() throws Exception {
-        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(1, 100)).toList();
+        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(1, 50)).toList();
         BulkInferenceRequestIterator requestIterator = requestIterator(requests);
         List<InferenceAction.Response> responses = Stream.generate(() -> mockInferenceResponse(RankedDocsResults.class))
             .limit(requests.size())
@@ -77,7 +77,7 @@ public void testSuccessfulExecution() throws Exception {
         doAnswer((invocation) -> {
             ActionListener<InferenceAction.Response> l = invocation.getArgument(1);
             if (randomBoolean()) {
-                Thread.sleep(between(0, 50));
+                Thread.sleep(between(0, 5));
             }
             l.onResponse(responses.get(requests.indexOf(invocation.getArgument(0, InferenceAction.Request.class))));
             return null;
@@ -101,7 +101,7 @@ public void testSuccessfulExecution() throws Exception {
             assertThat(output, hasSize(requests.size()));
             assertThat(output, contains(responses.stream().map(InferenceAction.Response::getResults).toArray()));
             verify(listener).onResponse(eq(output));
-        }, 60, TimeUnit.SECONDS);
+        });
     }
 
     @SuppressWarnings("unchecked")
@@ -126,14 +126,14 @@ public void testSuccessfulExecutionOnEmptyRequest() throws Exception {
 
     @SuppressWarnings("unchecked")
     public void testInferenceRunnerAlwaysFails() throws Exception {
-        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(1, 10)).toList();
+        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(1, 30)).toList();
         BulkInferenceRequestIterator requestIterator = requestIterator(requests);
 
         InferenceRunner inferenceRunner = mock(InferenceRunner.class);
         doAnswer(invocation -> {
             ActionListener<InferenceAction.Response> listener = invocation.getArgument(1);
             if (randomBoolean()) {
-                Thread.sleep(between(0, 500));
+                Thread.sleep(between(0, 5));
             }
             listener.onFailure(new RuntimeException("inference failure"));
             return null;
@@ -159,17 +159,17 @@ public void testInferenceRunnerAlwaysFails() throws Exception {
 
     @SuppressWarnings("unchecked")
     public void testInferenceRunnerSometimesFails() throws Exception {
-        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(2, 10)).toList();
+        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(2, 30)).toList();
         BulkInferenceRequestIterator requestIterator = requestIterator(requests);
 
         InferenceRunner inferenceRunner = mock(InferenceRunner.class);
         doAnswer(invocation -> {
             ActionListener<InferenceAction.Response> listener = invocation.getArgument(1);
             if (randomBoolean()) {
-                Thread.sleep(between(0, 500));
+                Thread.sleep(between(0, 5));
             }
 
-            if (requests.indexOf(invocation.getArgument(0, InferenceAction.Request.class)) % requests.size() == 0) {
+            if ((requests.indexOf(invocation.getArgument(0, InferenceAction.Request.class)) % requests.size()) == 0) {
                 listener.onFailure(new RuntimeException("inference failure"));
             } else {
                 listener.onResponse(mockInferenceResponse(RankedDocsResults.class));
@@ -199,7 +199,7 @@ public void testInferenceRunnerSometimesFails() throws Exception {
 
     @SuppressWarnings("unchecked")
     public void testBuildOutputFailure() throws Exception {
-        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(1, 10)).toList();
+        List<InferenceAction.Request> requests = Stream.generate(this::mockInferenceRequest).limit(between(1, 30)).toList();
         BulkInferenceRequestIterator requestIterator = requestIterator(requests);
 
         InferenceRunner inferenceRunner = mock(InferenceRunner.class);