Adding supported streaming tasks tests

Author: jonathan-buttner

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/RetryRule.java

@@ -20,13 +20,13 @@
 
 /**
  * Provides a way to retry a failed test. To use this functionality add something like the following to your test class:
- * </br>
+ * <p>
  * <code>
  * {@literal @}Rule
- * </br>
+ * <p>
  * public RetryRule retry = new RetryRule(3, TimeValue.timeValueSeconds(1));
  * </code>
- * </p>
+ * <p>
  * See {@link InferenceGetServicesIT#retry} for an example.
  */
 public class RetryRule implements TestRule {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java

@@ -57,6 +57,8 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.xpack.core.inference.results.ResultUtils.createInvalidChunkedResultException;
 import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
@@ -87,6 +89,7 @@ public class ElasticInferenceService extends SenderService {
     private EnumSet<TaskType> enabledTaskTypes;
     private final ModelRegistry modelRegistry;
     private final ElasticInferenceServiceAuthorizationHandler authorizationHandler;
+    private final CountDownLatch authorizationCompletedLatch = new CountDownLatch(1);
 
     public ElasticInferenceService(
         HttpRequestSender.Factory factory,
@@ -108,8 +111,12 @@ public ElasticInferenceService(
 
     private void getAuthorization() {
         try {
-            ActionListener<ElasticInferenceServiceAuthorization> listener = ActionListener.wrap(this::setEnabledTaskTypes, e -> {
+            ActionListener<ElasticInferenceServiceAuthorization> listener = ActionListener.wrap(result -> {
+                setEnabledTaskTypes(result);
+                authorizationCompletedLatch.countDown();
+            }, e -> {
                 // we don't need to do anything if there was a failure, everything is disabled by default
+                authorizationCompletedLatch.countDown();
             });
 
             authorizationHandler.getAuthorization(listener, getSender());
@@ -131,6 +138,17 @@ private static EnumSet<TaskType> filterTaskTypesByAuthorization(ElasticInference
         return implementedTaskTypes;
     }
 
+    // Default for testing
+    void waitForAuthorizationToComplete(TimeValue waitTime) {
+        try {
+            if (authorizationCompletedLatch.await(waitTime.getSeconds(), TimeUnit.SECONDS) == false) {
+                throw new IllegalStateException("The wait time has expired for authorization to complete.");
+            }
+        } catch (InterruptedException e) {
+            throw new IllegalStateException("Waiting for authorization to complete was interrupted");
+        }
+    }
+
     @Override
     public synchronized Set<TaskType> supportedStreamingTasks() {
         var enabledStreamingTaskTypes = EnumSet.of(TaskType.CHAT_COMPLETION);

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

@@ -739,6 +739,50 @@ public void testGetConfiguration_WithoutSupportedTaskTypes_WhenModelsReturnTaskO
         }
     }
 
+    public void testSupportedStreamingTasks_ReturnsChatCompletion_WhenAuthRespondsWithAValidModel() throws Exception {
+        String responseJson = """
+            {
+                "models": [
+                    {
+                      "model_name": "model-a",
+                      "task_types": ["embed/text/sparse", "chat"]
+                    }
+                ]
+            }
+            """;
+
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
+
+        var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+        try (var service = createServiceWithAuthHandler(senderFactory, getUrl(webServer))) {
+            service.waitForAuthorizationToComplete(TIMEOUT);
+            assertThat(service.supportedStreamingTasks(), is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.ANY)));
+            assertTrue(service.defaultConfigIds().isEmpty());
+        }
+    }
+
+    public void testSupportedStreamingTasks_ReturnsEmpty_WhenAuthRespondsWithoutChatCompletion() throws Exception {
+        String responseJson = """
+            {
+                "models": [
+                    {
+                      "model_name": "model-a",
+                      "task_types": ["embed/text/sparse"]
+                    }
+                ]
+            }
+            """;
+
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
+
+        var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+        try (var service = createServiceWithAuthHandler(senderFactory, getUrl(webServer))) {
+            service.waitForAuthorizationToComplete(TIMEOUT);
+            assertThat(service.supportedStreamingTasks(), is(EnumSet.noneOf(TaskType.class)));
+            assertTrue(service.defaultConfigIds().isEmpty());
+        }
+    }
+
     private ElasticInferenceService createServiceWithMockSender() {
         return createServiceWithMockSender(ElasticInferenceServiceAuthorizationTests.createEnabledAuth());
     }
@@ -788,4 +832,14 @@ private ElasticInferenceService createService(
             mockAuthHandler
         );
     }
+
+    private ElasticInferenceService createServiceWithAuthHandler(HttpRequestSender.Factory senderFactory, String eisGatewayUrl) {
+        return new ElasticInferenceService(
+            senderFactory,
+            createWithEmptySettings(threadPool),
+            new ElasticInferenceServiceComponents(eisGatewayUrl),
+            mockModelRegistry(),
+            new ElasticInferenceServiceAuthorizationHandler(eisGatewayUrl, threadPool)
+        );
+    }
 }