Adding tests

Author: jonathan-buttner

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

@@ -17,20 +17,19 @@
 
 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"]
-                    }
-                ]
-            }
-        """;
+                {
+                    "models": [
+                        {
+                            "model_name": "test_model",
+                            "task_types": ["embed/text/sparse", "chat"]
+                        }
+                    ]
+                }
+            """;
 
         try (var parser = createParser(JsonXContent.jsonXContent, json)) {
             var entity = ElasticInferenceServiceAuthorizationResponseEntity.PARSER.apply(parser, null);
@@ -47,13 +46,12 @@ public void testParseAllFields() throws IOException {
         }
     }
 
-
     public void testParsing_EmptyModels() throws IOException {
         String json = """
-            {
-                "models": []
-            }
-        """;
+                {
+                    "models": []
+                }
+            """;
 
         try (var parser = createParser(JsonXContent.jsonXContent, json)) {
             var entity = ElasticInferenceServiceAuthorizationResponseEntity.PARSER.apply(parser, null);

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

@@ -12,20 +12,25 @@
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.http.MockResponse;
 import org.elasticsearch.test.http.MockWebServer;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.core.inference.results.ChatCompletionResults;
 import org.elasticsearch.xpack.inference.external.http.HttpClientManager;
+import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests;
+import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 import org.junit.After;
 import org.junit.Before;
 import org.mockito.ArgumentCaptor;
 
 import java.io.IOException;
 import java.util.EnumSet;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool;
@@ -34,10 +39,12 @@
 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.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
 
 public class ElasticInferenceServiceAuthorizationHandlerTests extends ESTestCase {
     private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS);
@@ -185,6 +192,7 @@ public void testGetAuthorization_ReturnsAValidResponse() throws IOException {
             verifyNoMoreInteractions(logger);
         }
     }
+
     @SuppressWarnings("unchecked")
     public void testGetAuthorization_OnResponseCalledOnce() throws IllegalStateException {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
@@ -195,64 +203,63 @@ public void testGetAuthorization_OnResponseCalledOnce() throws IllegalStateExcep
         ActionListener<ElasticInferenceServiceAuthorization> listener = mock(ActionListener.class);
 
         String responseJson = """
-            {
-                "models": [
-                    {
-                      "model_name": "model-a",
-                      "task_types": ["embedding/text/sparse", "chat/completion"]
-                    }
-                ]
-            }
-        """;
+                {
+                    "models": [
+                        {
+                          "model_name": "model-a",
+                          "task_types": ["embed/text/sparse", "chat"]
+                        }
+                    ]
+                }
+            """;
 
         webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
 
         authHandler.getAuthorization(listener, senderFactory.createSender());
-        authHandler.waitForAuthRequestCompletion(TimeValue.timeValueSeconds(1));
+        authHandler.waitForAuthRequestCompletion(TIMEOUT);
 
         verify(listener, times(1)).onResponse(any());
+        var loggerArgsCaptor = ArgumentCaptor.forClass(String.class);
+        verify(logger, times(1)).debug(loggerArgsCaptor.capture());
+
+        var message = loggerArgsCaptor.getValue();
+        assertThat(message, is("Retrieving authorization information from the Elastic Inference Service."));
         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);
+    public void testGetAuthorization_InvalidResponse() throws IOException {
+        var senderMock = mock(Sender.class);
+        var senderFactory = mock(HttpRequestSender.Factory.class);
+        when(senderFactory.createSender()).thenReturn(senderMock);
 
-    try (var sender = senderFactory.createSender()) {
-        String responseJson = """
-                {
-                    "completion": [
-                        {
-                            "result": "some result 1"
-                        },
-                        {
-                            "result": "some result 2"
-                        }
-                    ]
-                }
-            """;
+        doAnswer(invocationOnMock -> {
+            ActionListener<InferenceServiceResults> listener = invocationOnMock.getArgument(4);
+            listener.onResponse(new ChatCompletionResults(List.of(new ChatCompletionResults.Result("awesome"))));
+            return Void.TYPE;
+        }).when(senderMock).sendWithoutQueuing(any(), any(), any(), any(), any());
 
-        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
+        var logger = mock(Logger.class);
+        var authHandler = new ElasticInferenceServiceAuthorizationHandler("abc", threadPool, logger);
+
+        try (var sender = senderFactory.createSender()) {
+            PlainActionFuture<ElasticInferenceServiceAuthorization> listener = new PlainActionFuture<>();
 
-        PlainActionFuture<ElasticInferenceServiceAuthorization> listener = new PlainActionFuture<>();
-        authHandler.getAuthorization(listener, sender);
+            authHandler.getAuthorization(listener, sender);
+            var result = listener.actionGet(TIMEOUT);
 
-        var authResponse = listener.actionGet(TIMEOUT);
-        assertTrue(authResponse.enabledTaskTypes().isEmpty());
-        assertFalse(authResponse.isEnabled());
+            assertThat(result, is(ElasticInferenceServiceAuthorization.newDisabledService()));
 
-        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"
-            ));
-    }
+            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: ChatCompletionResults"
+                )
+            );
+        }
 
     }