From 5bcbb92343424343f62465cf68e6ab00599f578f Mon Sep 17 00:00:00 2001
From: Keith Massey <keith.massey@elastic.co>
Date: Mon, 4 Dec 2023 10:25:40 -0600
Subject: [PATCH] fixing tests

---
 .../org/elasticsearch/http/SearchRestCancellationIT.java | 9 +++++----
 .../saml/TransportSamlInvalidateSessionActionTests.java  | 5 +++++
 .../elasticsearch/xpack/security/test/SecurityMocks.java | 4 +---
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/qa/smoke-test-http/src/javaRestTest/java/org/elasticsearch/http/SearchRestCancellationIT.java b/qa/smoke-test-http/src/javaRestTest/java/org/elasticsearch/http/SearchRestCancellationIT.java
index 73dd1525f8a08..8fbe9abe7f1b3 100644
--- a/qa/smoke-test-http/src/javaRestTest/java/org/elasticsearch/http/SearchRestCancellationIT.java
+++ b/qa/smoke-test-http/src/javaRestTest/java/org/elasticsearch/http/SearchRestCancellationIT.java
@@ -184,11 +184,12 @@ private static void ensureSearchTaskIsCancelled(String transportAction, Function
     private static void indexTestData() {
         for (int i = 0; i < 5; i++) {
             // Make sure we have a few segments
-            BulkRequestBuilder bulkRequestBuilder = client().prepareBulk().setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
-            for (int j = 0; j < 20; j++) {
-                bulkRequestBuilder.add(prepareIndex("test").setId(Integer.toString(i * 5 + j)).setSource("field", "value"));
+            try (BulkRequestBuilder bulkRequestBuilder = client().prepareBulk().setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)) {
+                for (int j = 0; j < 20; j++) {
+                    bulkRequestBuilder.add(prepareIndex("test").setId(Integer.toString(i * 5 + j)).setSource("field", "value"));
+                }
+                assertNoFailures(bulkRequestBuilder.get());
             }
-            assertNoFailures(bulkRequestBuilder.get());
         }
     }
 
diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/action/saml/TransportSamlInvalidateSessionActionTests.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/action/saml/TransportSamlInvalidateSessionActionTests.java
index 28f8a77c422f7..3958c8d629b4d 100644
--- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/action/saml/TransportSamlInvalidateSessionActionTests.java
+++ b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/action/saml/TransportSamlInvalidateSessionActionTests.java
@@ -259,6 +259,11 @@ protected <Request extends ActionRequest, Response extends ActionResponse> void
             ((Runnable) inv.getArguments()[1]).run();
             return null;
         }).when(securityIndex).prepareIndexIfNeededThenExecute(anyConsumer(), any(Runnable.class));
+        doAnswer(inv -> {
+            ((Runnable) inv.getArguments()[1]).run();
+            ((Runnable) inv.getArguments()[2]).run();
+            return null;
+        }).when(securityIndex).prepareIndexIfNeededThenExecute(anyConsumer(), any(Runnable.class), any(Runnable.class));
         doAnswer(inv -> {
             ((Runnable) inv.getArguments()[1]).run();
             return null;
diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/test/SecurityMocks.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/test/SecurityMocks.java
index c396e5a74dc27..995163572f888 100644
--- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/test/SecurityMocks.java
+++ b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/test/SecurityMocks.java
@@ -77,10 +77,8 @@ public static SecurityIndexManager mockSecurityIndexManager(String alias, boolea
         doAnswer(invocationOnMock -> {
             Runnable runnable = (Runnable) invocationOnMock.getArguments()[1];
             runnable.run();
-            Runnable onCompletion = (Runnable) invocationOnMock.getArguments()[2];
-            onCompletion.run();
             return null;
-        }).when(securityIndexManager).prepareIndexIfNeededThenExecute(anyConsumer(), any(Runnable.class), any(Runnable.class));
+        }).when(securityIndexManager).prepareIndexIfNeededThenExecute(anyConsumer(), any(Runnable.class));
         doAnswer(invocationOnMock -> {
             Runnable runnable = (Runnable) invocationOnMock.getArguments()[1];
             runnable.run();