Skip to content

Commit f8c8783

Browse files
[ML] Ensure EIS auth response exists before executing tests (elastic#128640)
* Add response before each test file for eis * Fixing the expected values * Adding response before cluster starts
1 parent ed141b1 commit f8c8783

File tree

4 files changed

+36
-24
lines changed

4 files changed

+36
-24
lines changed

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

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -24,17 +24,15 @@
2424

2525
public class BaseMockEISAuthServerTest extends ESRestTestCase {
2626

27-
// The reason we're retrying is there's a race condition between the node retrieving the
28-
// authorization response and running the test. Retrieving the authorization should be very fast since
29-
// we're hosting a local mock server but it's possible it could respond slower. So in the even of a test failure
30-
// we'll automatically retry after waiting a second.
31-
@Rule
32-
public RetryRule retry = new RetryRule(3, TimeValue.timeValueSeconds(1));
27+
protected static final MockElasticInferenceServiceAuthorizationServer mockEISServer =
28+
new MockElasticInferenceServiceAuthorizationServer();
3329

34-
private static final MockElasticInferenceServiceAuthorizationServer mockEISServer = MockElasticInferenceServiceAuthorizationServer
35-
.enabledWithRainbowSprinklesAndElser();
30+
static {
31+
// Ensure that the mock EIS server has an authorized response prior to the cluster starting
32+
mockEISServer.enqueueAuthorizeAllModelsResponse();
33+
}
3634

37-
private static final ElasticsearchCluster cluster = ElasticsearchCluster.local()
35+
private static ElasticsearchCluster cluster = ElasticsearchCluster.local()
3836
.distribution(DistributionType.DEFAULT)
3937
.setting("xpack.license.self_generated.type", "trial")
4038
.setting("xpack.security.enabled", "true")
@@ -52,9 +50,18 @@ public class BaseMockEISAuthServerTest extends ESRestTestCase {
5250

5351
// The reason we're doing this is to make sure the mock server is initialized first so we can get the address before communicating
5452
// it to the cluster as a setting.
53+
// Note: @ClassRule is executed once for the entire test class
5554
@ClassRule
5655
public static TestRule ruleChain = RuleChain.outerRule(mockEISServer).around(cluster);
5756

57+
// The reason we're retrying is there's a race condition between the node retrieving the
58+
// authorization response and running the test. Retrieving the authorization should be very fast since
59+
// we're hosting a local mock server but it's possible it could respond slower. So in the even of a test failure
60+
// we'll automatically retry after waiting a second.
61+
// Note: @Rule is executed for each test
62+
@Rule
63+
public RetryRule retry = new RetryRule(3, TimeValue.timeValueSeconds(1));
64+
5865
@Override
5966
protected String getTestRestCluster() {
6067
return cluster.getHttpAddresses();

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
package org.elasticsearch.xpack.inference;
1111

1212
import org.elasticsearch.inference.TaskType;
13+
import org.junit.BeforeClass;
1314

1415
import java.io.IOException;
1516
import java.util.List;
@@ -22,6 +23,12 @@
2223

2324
public class InferenceGetModelsWithElasticInferenceServiceIT extends BaseMockEISAuthServerTest {
2425

26+
@BeforeClass
27+
public static void init() {
28+
// Ensure the mock EIS server has an authorized response ready
29+
mockEISServer.enqueueAuthorizeAllModelsResponse();
30+
}
31+
2532
public void testGetDefaultEndpoints() throws IOException {
2633
var allModels = getAllModels();
2734
var chatCompletionModels = getModels("_all", TaskType.CHAT_COMPLETION);

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.elasticsearch.client.Request;
1313
import org.elasticsearch.common.Strings;
1414
import org.elasticsearch.inference.TaskType;
15+
import org.junit.BeforeClass;
1516

1617
import java.io.IOException;
1718
import java.util.List;
@@ -23,6 +24,12 @@
2324

2425
public class InferenceGetServicesIT extends BaseMockEISAuthServerTest {
2526

27+
@BeforeClass
28+
public static void init() {
29+
// Ensure the mock EIS server has an authorized response ready
30+
mockEISServer.enqueueAuthorizeAllModelsResponse();
31+
}
32+
2633
public void testGetServicesWithoutTaskType() throws IOException {
2734
List<Object> services = getAllServices();
2835
assertThat(services.size(), equalTo(23));

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

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,11 @@ public class MockElasticInferenceServiceAuthorizationServer implements TestRule
2626
public static MockElasticInferenceServiceAuthorizationServer enabledWithRainbowSprinklesAndElser() {
2727
var server = new MockElasticInferenceServiceAuthorizationServer();
2828

29+
server.enqueueAuthorizeAllModelsResponse();
30+
return server;
31+
}
32+
33+
public void enqueueAuthorizeAllModelsResponse() {
2934
String responseJson = """
3035
{
3136
"models": [
@@ -41,21 +46,7 @@ public static MockElasticInferenceServiceAuthorizationServer enabledWithRainbowS
4146
}
4247
""";
4348

44-
server.webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
45-
return server;
46-
}
47-
48-
public static MockElasticInferenceServiceAuthorizationServer disabled() {
49-
var server = new MockElasticInferenceServiceAuthorizationServer();
50-
51-
String responseJson = """
52-
{
53-
"models": []
54-
}
55-
""";
56-
57-
server.webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
58-
return server;
49+
webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
5950
}
6051

6152
public String getUrl() {

0 commit comments

Comments
 (0)