diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/BaseMockEISAuthServerTest.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/BaseMockEISAuthServerTest.java index 81c1a8dc7a5ba..4d6ed27edb9f9 100644 --- a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/BaseMockEISAuthServerTest.java +++ b/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(); diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetModelsWithElasticInferenceServiceIT.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetModelsWithElasticInferenceServiceIT.java index 42289c50864e6..8439dfb6af3ca 100644 --- a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetModelsWithElasticInferenceServiceIT.java +++ b/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); diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java index 9ed1b2a642f4c..91d4a4e37d6d5 100644 --- a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java +++ b/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 services = getAllServices(); assertThat(services.size(), equalTo(22)); diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockElasticInferenceServiceAuthorizationServer.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockElasticInferenceServiceAuthorizationServer.java index 3ea011c1317cc..cf798d5c94364 100644 --- a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockElasticInferenceServiceAuthorizationServer.java +++ b/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() {