Fixing tests

Author: jonathan-buttner

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/SenderServiceTests.java

@@ -116,10 +116,21 @@ public void testReturnsValidationException_WhenQueryIsNullForRerankTaskType() th
             var model = mock(Model.class);
             when(model.getTaskType()).thenReturn(TaskType.RERANK);
 
-            PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
-
-            testService.infer(model, null, null, null, List.of("test input"), false, Map.of(), InputType.SEARCH, null, listener);
-            var exception = expectThrows(ValidationException.class, () -> listener.actionGet(TIMEOUT));
+            var exception = expectThrows(
+                ValidationException.class,
+                () -> testService.infer(
+                    model,
+                    null,
+                    null,
+                    null,
+                    List.of("test input"),
+                    false,
+                    Map.of(),
+                    InputType.SEARCH,
+                    null,
+                    new PlainActionFuture<>()
+                )
+            );
 
             assertThat(exception.getMessage(), containsString("Rerank task type requires a non-null query field"));
         }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java

@@ -899,23 +899,23 @@ public void testInfer_ReturnsValidationException_WhenSendingRerankRequest_Withou
                 new DefaultSecretSettings(new SecureString(secret.toCharArray()))
             );
 
-            PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
-            service.infer(
-                model,
-                // null query string will trigger validation error
-                null,
-                null,
-                null,
-                List.of("abc"),
-                false,
-                new HashMap<>(),
-                InputType.INGEST,
-                InferenceAction.Request.DEFAULT_TIMEOUT,
-                listener
+            var exception = expectThrows(
+                ValidationException.class,
+                () -> service.infer(
+                    model,
+                    // null query string will trigger validation error
+                    null,
+                    null,
+                    null,
+                    List.of("abc"),
+                    false,
+                    new HashMap<>(),
+                    InputType.INGEST,
+                    InferenceAction.Request.DEFAULT_TIMEOUT,
+                    new PlainActionFuture<>()
+                )
             );
 
-            var exception = expectThrows(ValidationException.class, () -> listener.actionGet(TIMEOUT));
-
             assertThat(exception.getMessage(), containsString("Rerank task type requires a non-null query field"));
         }
     }