[ML] Ensure EIS auth response exists before executing tests

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

@@ -23,17 +23,15 @@
 
 public class BaseMockEISAuthServerTest extends ESRestTestCase {
 
-    // The reason we're retrying is there's a race condition between the node retrieving the
-    // 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.
-    @Rule
-    public RetryRule retry = new RetryRule(3, TimeValue.timeValueSeconds(1));
+    protected static final MockElasticInferenceServiceAuthorizationServer mockEISServer =
+        new MockElasticInferenceServiceAuthorizationServer();
 
-    private static final MockElasticInferenceServiceAuthorizationServer mockEISServer = MockElasticInferenceServiceAuthorizationServer
-        .enabledWithRainbowSprinklesAndElser();
+    static {
+        // Ensure that the mock EIS server has an authorized response prior to the cluster starting
+        mockEISServer.enqueueAuthorizeAllModelsResponse();
+    }
 
-    private static final ElasticsearchCluster cluster = ElasticsearchCluster.local()
+    private static ElasticsearchCluster cluster = ElasticsearchCluster.local()
         .distribution(DistributionType.DEFAULT)
         .setting("xpack.license.self_generated.type", "trial")
         .setting("xpack.security.enabled", "true")
@@ -50,9 +48,18 @@ 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);
 
+    // The reason we're retrying is there's a race condition between the node retrieving the
+    // 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));
+
     @Override
     protected String getTestRestCluster() {
         return cluster.getHttpAddresses();

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));

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() {