Add response before each test file for eis

Author: jonathan-buttner

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

@@ -27,11 +27,12 @@ public class BaseMockEISAuthServerTest extends ESRestTestCase {
     // authorization response and running the test. Retrieving the authorization should be very fast since
     // we're hosting a local mock server but it's possible it could respond slower. So in the even of a test failure
     // we'll automatically retry after waiting a second.
+    // Note: @Rule is executed for each test
     @Rule
     public RetryRule retry = new RetryRule(3, TimeValue.timeValueSeconds(1));
 
-    private static final MockElasticInferenceServiceAuthorizationServer mockEISServer = MockElasticInferenceServiceAuthorizationServer
-        .enabledWithRainbowSprinklesAndElser();
+    protected static final MockElasticInferenceServiceAuthorizationServer mockEISServer =
+        new MockElasticInferenceServiceAuthorizationServer();
 
     private static final ElasticsearchCluster cluster = ElasticsearchCluster.local()
         .distribution(DistributionType.DEFAULT)
@@ -50,6 +51,7 @@ public class BaseMockEISAuthServerTest extends ESRestTestCase {
 
     // The reason we're doing this is to make sure the mock server is initialized first so we can get the address before communicating
     // it to the cluster as a setting.
+    // Note: @ClassRule is executed once for the entire test class
     @ClassRule
     public static TestRule ruleChain = RuleChain.outerRule(mockEISServer).around(cluster);
 

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

@@ -10,6 +10,7 @@
 package org.elasticsearch.xpack.inference;
 
 import org.elasticsearch.inference.TaskType;
+import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.List;
@@ -22,6 +23,12 @@
 
 public class InferenceGetModelsWithElasticInferenceServiceIT extends BaseMockEISAuthServerTest {
 
+    @BeforeClass
+    public static void init() {
+        // Ensure the mock EIS server has an authorized response ready
+        mockEISServer.enqueueAuthorizeAllModelsResponse();
+    }
+
     public void testGetDefaultEndpoints() throws IOException {
         var allModels = getAllModels();
         var chatCompletionModels = getModels("_all", TaskType.CHAT_COMPLETION);

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.client.Request;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.TaskType;
+import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.List;
@@ -23,6 +24,12 @@
 
 public class InferenceGetServicesIT extends BaseMockEISAuthServerTest {
 
+    @BeforeClass
+    public static void init() {
+        // Ensure the mock EIS server has an authorized response ready
+        mockEISServer.enqueueAuthorizeAllModelsResponse();
+    }
+
     public void testGetServicesWithoutTaskType() throws IOException {
         List<Object> services = getAllServices();
         assertThat(services.size(), equalTo(22));
@@ -101,7 +108,7 @@ public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
 
     public void testGetServicesWithRerankTaskType() throws IOException {
         List<Object> services = getServices(TaskType.RERANK);
-        assertThat(services.size(), equalTo(8));
+        assertThat(services.size(), equalTo(9));
 
         var providers = providers(services);
 
@@ -124,7 +131,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
 
     public void testGetServicesWithCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.COMPLETION);
-        assertThat(services.size(), equalTo(12));
+        assertThat(services.size(), equalTo(13));
 
         var providers = providers(services);
 

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

@@ -26,6 +26,11 @@ public class MockElasticInferenceServiceAuthorizationServer implements TestRule
     public static MockElasticInferenceServiceAuthorizationServer enabledWithRainbowSprinklesAndElser() {
         var server = new MockElasticInferenceServiceAuthorizationServer();
 
+        server.enqueueAuthorizeAllModelsResponse();
+        return server;
+    }
+
+    public void enqueueAuthorizeAllModelsResponse() {
         String responseJson = """
             {
                 "models": [
@@ -41,21 +46,7 @@ public static MockElasticInferenceServiceAuthorizationServer enabledWithRainbowS
             }
             """;
 
-        server.webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
-        return server;
-    }
-
-    public static MockElasticInferenceServiceAuthorizationServer disabled() {
-        var server = new MockElasticInferenceServiceAuthorizationServer();
-
-        String responseJson = """
-            {
-                "models": []
-            }
-            """;
-
-        server.webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
-        return server;
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
     }
 
     public String getUrl() {