add AuthHandler, AuthRequest, and AuthResponseEntity tests

Author: brendan-jugan-elastic

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/elastic/ElasticInferenceServiceAuthorizationResponseEntity.java

@@ -163,4 +163,17 @@ public List<? extends InferenceResults> transformToLegacyFormat() {
     public Map<String, Object> asMap() {
         throw new UnsupportedOperationException("Not implemented");
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ElasticInferenceServiceAuthorizationResponseEntity that = (ElasticInferenceServiceAuthorizationResponseEntity) o;
+        return Objects.equals(authorizedModels, that.authorizedModels);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(authorizedModels);
+    }
 }

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -25,6 +26,8 @@
 
 import java.util.Locale;
 import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.xpack.core.inference.action.InferenceAction.Request.DEFAULT_TIMEOUT;
 import static org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService.ELASTIC_INFERENCE_SERVICE_IDENTIFIER;
@@ -48,6 +51,7 @@ private static ResponseHandler createAuthResponseHandler() {
     private final String baseUrl;
     private final ThreadPool threadPool;
     private final Logger logger;
+    private final CountDownLatch requestCompleteLatch = new CountDownLatch(1);
 
     public ElasticInferenceServiceAuthorizationHandler(@Nullable String baseUrl, ThreadPool threadPool) {
         this.baseUrl = baseUrl;
@@ -92,6 +96,7 @@ public void getAuthorization(ActionListener<ElasticInferenceServiceAuthorization
                     );
                     listener.onResponse(ElasticInferenceServiceAuthorization.newDisabledService());
                 }
+                requestCompleteLatch.countDown();
             }, e -> {
                 Throwable exception = e;
                 if (e instanceof ElasticsearchWrapperException wrapperException) {
@@ -100,6 +105,7 @@ public void getAuthorization(ActionListener<ElasticInferenceServiceAuthorization
 
                 logger.warn(Strings.format(FAILED_TO_RETRIEVE_MESSAGE + " Encountered an exception: %s", exception));
                 listener.onResponse(ElasticInferenceServiceAuthorization.newDisabledService());
+                requestCompleteLatch.countDown();
             });
 
             var request = new ElasticInferenceServiceAuthorizationRequest(baseUrl, getCurrentTraceInfo());
@@ -117,4 +123,14 @@ private TraceContext getCurrentTraceInfo() {
 
         return new TraceContext(traceParent, traceState);
     }
+
+    void waitForAuthRequestCompletion(TimeValue timeValue) throws IllegalStateException {
+        try {
+            if (requestCompleteLatch.await(timeValue.getMillis(), TimeUnit.MILLISECONDS) == false) {
+                throw new IllegalStateException("The authorization request did not complete.");
+            }
+        } catch (IllegalStateException | InterruptedException e) {
+            logger.warn("Interrupted while waiting for the authorization request to complete", e);
+        }
+    }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/elastic/ElasticInferenceServiceAuthorizationRequestTests.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.request.elastic;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.inference.telemetry.TraceContext;
+import org.junit.Before;
+
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.is;
+
+public class ElasticInferenceServiceAuthorizationRequestTests extends ESTestCase {
+
+    private TraceContext traceContext;
+
+    @Before
+    public void init() {
+        traceContext = new TraceContext("dummyTraceParent", "dummyTraceState");
+    }
+
+    public void testCreateUriThrowsForInvalidBaseUrl() {
+        String invalidUrl = "http://invalid-url^";
+
+        ElasticsearchStatusException exception = assertThrows(
+            ElasticsearchStatusException.class,
+            () -> new ElasticInferenceServiceAuthorizationRequest(invalidUrl, traceContext)
+        );
+
+        assertThat(exception.status(), is(RestStatus.BAD_REQUEST));
+        assertThat(exception.getMessage(), containsString("Failed to create URI for service"));
+    }
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/response/elastic/ElasticInferenceServiceAuthorizationResponseEntityTests.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.response.elastic;
+
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.json.JsonXContent;
+
+import java.io.IOException;
+import java.util.EnumSet;
+import java.util.List;
+
+import static org.hamcrest.Matchers.is;
+
+
+public class ElasticInferenceServiceAuthorizationResponseEntityTests extends ESTestCase {
+
+    public void testParseAllFields() throws IOException {
+        String json = """
+            {
+                "models": [
+                    {
+                        "model_name": "test_model",
+                        "task_types": ["embedding/text/sparse", "chat/completion"]
+                    }
+                ]
+            }
+        """;
+
+        try (var parser = createParser(JsonXContent.jsonXContent, json)) {
+            var entity = ElasticInferenceServiceAuthorizationResponseEntity.PARSER.apply(parser, null);
+            var expected = new ElasticInferenceServiceAuthorizationResponseEntity(
+                List.of(
+                    new ElasticInferenceServiceAuthorizationResponseEntity.AuthorizedModel(
+                        "test_model",
+                        EnumSet.of(TaskType.SPARSE_EMBEDDING, TaskType.CHAT_COMPLETION)
+                    )
+                )
+            );
+
+            assertThat(entity, is(expected));
+        }
+    }
+
+
+    public void testParsing_EmptyModels() throws IOException {
+        String json = """
+            {
+                "models": []
+            }
+        """;
+
+        try (var parser = createParser(JsonXContent.jsonXContent, json)) {
+            var entity = ElasticInferenceServiceAuthorizationResponseEntity.PARSER.apply(parser, null);
+            var expected = new ElasticInferenceServiceAuthorizationResponseEntity(List.of());
+
+            assertThat(entity, is(expected));
+        }
+    }
+
+}

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

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.inference.services.elastic.authorization;
 
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
@@ -32,6 +33,7 @@
 import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
 import static org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender.MAX_RETIES;
 import static org.hamcrest.Matchers.is;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -183,4 +185,75 @@ public void testGetAuthorization_ReturnsAValidResponse() throws IOException {
             verifyNoMoreInteractions(logger);
         }
     }
+    @SuppressWarnings("unchecked")
+    public void testGetAuthorization_OnResponseCalledOnce() throws IllegalStateException {
+        var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+        var eisGatewayUrl = getUrl(webServer);
+        var logger = mock(Logger.class);
+        var authHandler = new ElasticInferenceServiceAuthorizationHandler(eisGatewayUrl, threadPool, logger);
+
+        ActionListener<ElasticInferenceServiceAuthorization> listener = mock(ActionListener.class);
+
+        String responseJson = """
+            {
+                "models": [
+                    {
+                      "model_name": "model-a",
+                      "task_types": ["embedding/text/sparse", "chat/completion"]
+                    }
+                ]
+            }
+        """;
+
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
+
+        authHandler.getAuthorization(listener, senderFactory.createSender());
+        authHandler.waitForAuthRequestCompletion(TimeValue.timeValueSeconds(1));
+
+        verify(listener, times(1)).onResponse(any());
+        verifyNoMoreInteractions(logger);
+    }
+
+public void testGetAuthorization_InvalidResponse() throws IOException {
+    var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+    var eisGatewayUrl = getUrl(webServer);
+    var logger = mock(Logger.class);
+    var authHandler = new ElasticInferenceServiceAuthorizationHandler(eisGatewayUrl, threadPool, logger);
+
+    try (var sender = senderFactory.createSender()) {
+        String responseJson = """
+                {
+                    "completion": [
+                        {
+                            "result": "some result 1"
+                        },
+                        {
+                            "result": "some result 2"
+                        }
+                    ]
+                }
+            """;
+
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
+
+        PlainActionFuture<ElasticInferenceServiceAuthorization> listener = new PlainActionFuture<>();
+        authHandler.getAuthorization(listener, sender);
+
+        var authResponse = listener.actionGet(TIMEOUT);
+        assertTrue(authResponse.enabledTaskTypes().isEmpty());
+        assertFalse(authResponse.isEnabled());
+
+        var loggerArgsCaptor = ArgumentCaptor.forClass(String.class);
+        verify(logger).warn(loggerArgsCaptor.capture());
+        var message = loggerArgsCaptor.getValue();
+        assertThat(
+            message,
+            is(
+                "Failed to retrieve the authorization information from the Elastic Inference Service."
+                    + " Received an invalid response type: InferenceServiceResults"
+            ));
+    }
+
+    }
+
 }