Fixing tests

Author: jonathan-buttner

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

@@ -27,7 +27,6 @@
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
-import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceComponents;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMFeature;
 import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMService;
@@ -61,7 +60,6 @@ public class AuthorizationPoller extends AllocatedPersistentTask {
     private final AtomicBoolean shutdown = new AtomicBoolean(false);
     private final ElasticInferenceServiceSettings elasticInferenceServiceSettings;
     private final AtomicBoolean initialized = new AtomicBoolean(false);
-    private final ElasticInferenceServiceComponents elasticInferenceServiceComponents;
     private final Client client;
     private final CountDownLatch receivedFirstAuthResponseLatch = new CountDownLatch(1);
     private final CCMFeature ccmFeature;
@@ -118,9 +116,6 @@ private AuthorizationPoller(TaskFields taskFields, Parameters parameters) {
         this.authorizationHandler = Objects.requireNonNull(authorizationRequestHandler);
         this.sender = Objects.requireNonNull(sender);
         this.elasticInferenceServiceSettings = Objects.requireNonNull(elasticInferenceServiceSettings);
-        this.elasticInferenceServiceComponents = new ElasticInferenceServiceComponents(
-            elasticInferenceServiceSettings.getElasticInferenceServiceUrl()
-        );
         this.modelRegistry = Objects.requireNonNull(modelRegistry);
         this.client = new OriginSettingClient(Objects.requireNonNull(client), ClientHelper.INFERENCE_ORIGIN);
         this.ccmFeature = Objects.requireNonNull(ccmFeature);

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

@@ -44,6 +44,7 @@
 import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMFeatureTests.createMockCCMFeature;
 import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMServiceTests.createMockCCMService;
 import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.createAuthorizedEndpoint;
+import static org.elasticsearch.xpack.inference.services.elastic.response.AuthorizationResponseEntityTests.createInvalidTaskTypeAuthorizedEndpoint;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -336,7 +337,7 @@ public void testSendsAuthorizationRequest_ButDoesNotStoreAnyModels_WhenTheirInfe
 
     public void testDoesNotAttemptToStoreModelIds_ThatHaveATaskTypeThatTheEISIntegration_DoesNotSupport() {
         var url = "eis-url";
-        var completionModel = createAuthorizedEndpoint(TaskType.COMPLETION);
+        var invalidTaskTypeEndpoint = createInvalidTaskTypeAuthorizedEndpoint();
 
         var mockRegistry = mock(ModelRegistry.class);
         when(mockRegistry.isReady()).thenReturn(true);
@@ -345,7 +346,7 @@ public void testDoesNotAttemptToStoreModelIds_ThatHaveATaskTypeThatTheEISIntegra
         var mockAuthHandler = mock(ElasticInferenceServiceAuthorizationRequestHandler.class);
         doAnswer(invocation -> {
             ActionListener<AuthorizationModel> listener = invocation.getArgument(0);
-            listener.onResponse(AuthorizationModel.of(new AuthorizationResponseEntity(List.of(completionModel)), url));
+            listener.onResponse(AuthorizationModel.of(new AuthorizationResponseEntity(List.of(invalidTaskTypeEndpoint)), url));
             return Void.TYPE;
         }).when(mockAuthHandler).getAuthorization(any(), any());
 

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

@@ -203,7 +203,15 @@ public void testGetAuthorization_ReturnsAValidResponse() throws IOException {
             var authResponse = listener.actionGet(TIMEOUT);
             assertThat(
                 authResponse.getTaskTypes(),
-                is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING, TaskType.TEXT_EMBEDDING, TaskType.RERANK))
+                is(
+                    EnumSet.of(
+                        TaskType.CHAT_COMPLETION,
+                        TaskType.SPARSE_EMBEDDING,
+                        TaskType.TEXT_EMBEDDING,
+                        TaskType.RERANK,
+                        TaskType.COMPLETION
+                    )
+                )
             );
             assertThat(authResponse.getEndpointIds(), containsInAnyOrder(responseData.inferenceIds().toArray(String[]::new)));
             assertTrue(authResponse.isAuthorized());

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

@@ -493,6 +493,23 @@ public static AuthorizationResponseEntity createResponse() {
         );
     }
 
+    public static AuthorizationResponseEntity.AuthorizedEndpoint createInvalidTaskTypeAuthorizedEndpoint() {
+        var id = randomAlphaOfLength(10);
+        var name = randomAlphaOfLength(10);
+        var status = randomFrom("ga", "beta", "preview");
+
+        return new AuthorizationResponseEntity.AuthorizedEndpoint(
+            id,
+            name,
+            createTaskTypeObject("invalid/task/type", TaskType.ANY.toString()),
+            status,
+            null,
+            "",
+            "",
+            null
+        );
+    }
+
     public static AuthorizationResponseEntity.AuthorizedEndpoint createAuthorizedEndpoint(TaskType taskType) {
         var id = randomAlphaOfLength(10);
         var name = randomAlphaOfLength(10);
@@ -572,7 +589,15 @@ public void testParseAllFields() throws IOException {
 
             assertThat(
                 authModel.getTaskTypes(),
-                is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING, TaskType.TEXT_EMBEDDING, TaskType.RERANK))
+                is(
+                    EnumSet.of(
+                        TaskType.CHAT_COMPLETION,
+                        TaskType.SPARSE_EMBEDDING,
+                        TaskType.TEXT_EMBEDDING,
+                        TaskType.RERANK,
+                        TaskType.COMPLETION
+                    )
+                )
             );
             assertThat(
                 authModel.getEndpoints(responseData.inferenceIds()),